diff --git a/.gitignore b/.gitignore index cfd9ffb..2c19cf6 100644 --- a/.gitignore +++ b/.gitignore @@ -1,19 +1,13 @@ # Python *.py[cod] -*$py.class __pycache__/ *.so *.egg *.egg-info/ dist/ build/ -.Python -env/ venv/ -venv311/ -.venv -pip-log.txt -pip-delete-this-directory.txt +.venv/ # IDE .idea/ @@ -28,32 +22,25 @@ pip-delete-this-directory.txt # Project specific - Generated files *.log -training*.log -nohup.out # Project specific - Visualizations *.png *.jpg -umap*.png -# Project specific - Data -data/ -artifacts/ -wordnet/*.csv +# Project specific - Data (only root, not src/taxembed/data/) +/data/ +/artifacts/ -# Project specific - Compiled extensions -hype/adjacency_matrix_dataset.cpp -hype/graph_dataset.cpp -*.so +# Examples output +examples/*.pth +examples/*.png # macOS .DS_Store -# Testing +# Testing & Code Quality .pytest_cache/ .coverage htmlcov/ -.tox/ - -# Ruff .ruff_cache/ +.mypy_cache/ diff --git a/docs/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md similarity index 100% rename from docs/CODE_OF_CONDUCT.md rename to CODE_OF_CONDUCT.md diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..9a605c2 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,163 @@ +# Contributing to taxembed + +We want to make contributing to this project as easy and transparent as possible. + +## Development Setup + +1. Clone the repository: + +```bash +git clone https://github.com/jcoludar/taxembed.git +cd taxembed +``` + +2. Install development dependencies using uv: + +```bash +uv sync +``` + +## Pull Requests + +We actively welcome your pull requests. + +1. Fork the repo and create your branch from `main`. +2. If you've added code that should be tested, add tests in the `tests/` directory. +3. If you've changed APIs, update the documentation. +4. Ensure the test suite passes: `uv run pytest` +5. Make sure your code passes linting: `uv run ruff check src/ scripts/` +6. Format your code: `uv run ruff format src/ scripts/` + +## Code Quality + +This project maintains high code quality standards using modern Python tools: + +- **Ruff**: Fast linting and formatting +- **MyPy**: Static type checking +- **Pytest**: Comprehensive test suite + +### Linting with Ruff + +Check for linting issues: + +```bash +uv run ruff check . +``` + +Fix linting issues automatically: + +```bash +# Safe fixes only +uv run ruff check --fix . + +# Include unsafe fixes (e.g., unused imports) +uv run ruff check --fix --unsafe-fixes . +``` + +View detailed error explanations: + +```bash +uv run ruff check --output-format=full . +``` + +### Formatting with Ruff + +Format code to match project style: + +```bash +# Format all Python files +uv run ruff format . + +# Check formatting without making changes +uv run ruff format --check . +``` + +### Type Checking with MyPy + +Run static type analysis: + +```bash +# Check all source files +uv run mypy src/taxembed + +# Check specific module +uv run mypy src/taxembed/models/ + +# Show more detailed error messages +uv run mypy --show-error-codes src/taxembed +``` + +**Note:** This project uses gradual typing. Most modules are currently exempt from strict type checking (see `pyproject.toml`). When adding type hints to a module, remove it from the `[[tool.mypy.overrides]]` section. + +### Testing + +Run the test suite: + +```bash +# Run all tests +uv run pytest + +# Run with coverage report +uv run pytest --cov=src/taxembed --cov-report=term-missing + +# Run specific test file +uv run pytest tests/test_models.py + +# Run with verbose output +uv run pytest -v +``` + +### Complete Quality Check + +Run all quality checks at once: + +```bash +# Lint, format check, type check, and test +uv run ruff check . && \ +uv run ruff format --check . && \ +uv run mypy src/taxembed && \ +uv run pytest +``` + +## Coding Style + +- **Follow PEP 8 guidelines** (enforced by Ruff) +- **Use type hints** for all new code (function signatures and return types) +- **Write docstrings** for all public functions and classes (Google style preferred) +- **Keep lines under 100 characters** (enforced by Ruff formatter) +- **Use meaningful variable and function names** (avoid single letters except in loops) +- **Prefer explicit over implicit** (e.g., `zip(a, b, strict=True)`) + +### Type Hints Guidelines + +```python +# Good: Complete type hints +def process_data( + input_path: Path, + batch_size: int = 32, + verbose: bool = False, +) -> dict[str, np.ndarray]: + """Process data from file. + + Args: + input_path: Path to input file + batch_size: Number of items per batch + verbose: Enable verbose logging + + Returns: + Dictionary mapping names to arrays + """ + ... + +# Bad: Missing type hints +def process_data(input_path, batch_size=32, verbose=False): + ... +``` + +## Issues + +We use GitHub issues to track public bugs. Please ensure your description is clear and has sufficient instructions to reproduce the issue. + +## License + +By contributing to taxembed, you agree that your contributions will be licensed under the CC-BY-NC 4.0 license found in the LICENSE file in the root directory of this source tree. diff --git a/Makefile b/Makefile deleted file mode 100644 index 2783027..0000000 --- a/Makefile +++ /dev/null @@ -1,65 +0,0 @@ -.PHONY: help install lint format test clean train check - -help: - @echo "taxembed - Hierarchical Poincaré Embeddings for Taxonomy" - @echo "" - @echo "Available commands:" - @echo " make install Install dependencies with uv" - @echo " make lint Check code with ruff" - @echo " make format Format code with ruff" - @echo " make test Run tests with pytest" - @echo " make train Train small model (quick test)" - @echo " make check Run sanity checks" - @echo " make clean Remove build artifacts" - @echo "" - @echo "📖 See docs/ for detailed guides" - -install: - @echo "Installing dependencies with uv..." - uv sync - @echo "✅ Installation complete" - -install-dev: - @echo "Installing with dev dependencies..." - uv sync --all-extras - @echo "✅ Dev installation complete" - -lint: - uv run ruff check src/ scripts/ - -lint-fix: - uv run ruff check --fix src/ scripts/ - -format: - uv run ruff format src/ scripts/ - -test: - uv run pytest - -test-cov: - uv run pytest --cov=src/taxembed --cov-report=html - -train: - @echo "Training small model for 1 epoch (sanity check)..." - python train_small.py --epochs 1 - @echo "✅ Quick training test complete" - -check: - @echo "Running sanity checks..." - python final_sanity_check.py - @echo "✅ Sanity checks passed" - -clean: - @echo "Cleaning build artifacts..." - rm -rf build/ - rm -rf dist/ - rm -rf *.egg-info - rm -rf .pytest_cache/ - rm -rf .ruff_cache/ - rm -rf htmlcov/ - find . -type d -name __pycache__ -exec rm -rf {} + - find . -type f -name "*.pyc" -delete - find . -type f -name "*.so" -delete - @echo "✅ Cleanup complete" - -.DEFAULT_GOAL := help diff --git a/QUICKSTART.md b/QUICKSTART.md deleted file mode 100644 index 3cfbb0b..0000000 --- a/QUICKSTART.md +++ /dev/null @@ -1,151 +0,0 @@ -# Quick Start Guide - -Get started with hierarchical Poincaré taxonomy embeddings in minutes. - -## Prerequisites - -- Python 3.11 or higher -- [uv](https://github.com/astral-sh/uv) package manager (recommended) - -## Installation - -1. **Clone the repository:** -```bash -git clone https://github.com/jcoludar/taxembed.git -cd poincare-embeddings -``` - -2. **Install dependencies:** -```bash -make install -# or: uv sync -``` - -That's it! No compilation needed. - -## Use Pre-trained Model - -A production-ready model is included in `small_model_28epoch/`: - -```python -import torch -import pandas as pd - -# Load embeddings -ckpt = torch.load('small_model_28epoch/taxonomy_model_small_best.pth') -embeddings = ckpt['embeddings'] # 92,290 organisms × 10 dimensions - -# Load TaxID mapping -mapping = pd.read_csv('data/taxonomy_edges_small.mapping.tsv', - sep='\t', header=None, names=['idx', 'taxid']) -``` - -## Train New Model - -### Step 1: Prepare Data (if needed) - -Download NCBI taxonomy: -```bash -taxembed-download -# or: python prepare_taxonomy_data.py -``` - -Build transitive closure (975K training pairs): -```bash -taxembed-prepare -# or: python build_transitive_closure.py -``` - -### Step 2: Train - -```bash -taxembed-train -# or: python train_small.py -``` - -Training takes ~2.5 hours on M3 Mac CPU. The script includes: -- Real-time progress bars -- Early stopping (patience=5) -- Automatic best model saving - -### Step 3: Visualize - -```bash -taxembed-visualize small_model_28epoch/taxonomy_model_small_best.pth -# or: python visualize_multi_groups.py small_model_28epoch/taxonomy_model_small_best.pth -``` - -Generates UMAP visualization with key taxonomic groups highlighted. - -### Step 4: Verify - -```bash -taxembed-check -# or: python final_sanity_check.py -``` - -Runs comprehensive validation of models and data files. - -## Development - -### Check Code Quality - -```bash -make lint # Check with ruff -make format # Format code -make test # Run tests -``` - -### Quick Sanity Check - -```bash -make check # Run final_sanity_check.py -make train # Train for 1 epoch (test) -``` - -## Common Commands - -| Task | Command | -|------|---------| -| Install dependencies | `make install` | -| Train (1 epoch test) | `make train` | -| Check code quality | `make lint` | -| Format code | `make format` | -| Run tests | `make test` | -| Sanity check | `make check` | -| Clean artifacts | `make clean` | -| Show help | `make help` | - -## Troubleshooting - -### Missing data files - -If `train_small.py` fails with "Training data not found": -```bash -python build_transitive_closure.py -``` - -### Module not found - -Make sure you've installed dependencies: -```bash -make install -``` - -Or activate the virtual environment: -```bash -source venv311/bin/activate # if using venv -``` - -## Next Steps - -- See **docs/** for detailed guides: - - `docs/TRAIN_SMALL_GUIDE.md` - Training documentation - - `docs/JOURNEY.md` - Development history - - `docs/FINAL_STATUS.md` - Production status -- Review **README.md** for architecture details -- Check **small_model_28epoch/** for production model - -## Support - -For questions, see the documentation in **docs/** or open a GitHub issue. diff --git a/README.md b/README.md index 3827bdc..2bce3ff 100644 --- a/README.md +++ b/README.md @@ -6,8 +6,8 @@ **Learn hierarchical embeddings of NCBI's biological taxonomy in hyperbolic space.** -✅ **Production model included:** 92K organisms, loss 0.472, epoch 28 -📊 **Validated:** 100% ball constraint compliance, clear hierarchical clustering +✅ **Production model included:** 92K organisms, loss 0.472, epoch 28 +📊 **Validated:** 100% ball constraint compliance, clear hierarchical clustering 📁 **Location:** `small_model_28epoch/` This project extends Facebook Research's Poincaré embeddings with hierarchical features specifically designed for deep taxonomic hierarchies (38 levels, 2.7M organisms). @@ -31,6 +31,7 @@ This project extends Facebook Research's Poincaré embeddings with hierarchical ### **Production Model Available** ⭐ A pre-trained model is included in `small_model_28epoch/`: + - **92,290 organisms** embedded in 10 dimensions - **Best epoch:** 28, **Loss:** 0.472 - **100% ball constraint** compliance @@ -41,19 +42,15 @@ A pre-trained model is included in `small_model_28epoch/`: ```bash # Clone the repository git clone https://github.com/jcoludar/taxembed.git -cd poincare-embeddings +cd taxembed -# Install with uv (recommended) -make install -# or: uv sync +# Install with uv +uv sync -# Alternative: pip -python3.11 -m venv venv311 -source venv311/bin/activate -pip install -r requirements.txt ``` After installation, the unified CLI is available: + - `taxembed train -as ` - Train model for any clade (auto-builds dataset) - `taxembed visualize ` - Visualize results with automatic best checkpoint - `taxembed-download` - Download NCBI taxonomy (legacy, auto-handled by train) @@ -73,13 +70,14 @@ ckpt = torch.load('small_model_28epoch/taxonomy_model_small_best.pth') embeddings = ckpt['embeddings'] # Shape: (92290, 10) # Load TaxID mapping -mapping = pd.read_csv('data/taxonomy_edges_small.mapping.tsv', +mapping = pd.read_csv('data/taxonomy_edges_small.mapping.tsv', sep='\t', header=None, names=['idx', 'taxid']) ``` ### **Train New Model** **Using unified CLI** (recommended - easiest): + ```bash # Train any clade by name or TaxID (auto-builds dataset, downloads taxonomy if needed) taxembed train Cnidaria -as cnidaria --epochs 100 --lambda 0.1 @@ -93,6 +91,7 @@ taxembed visualize echinoderms --children 1 # Color by grandchildren ``` **Using legacy CLI commands**: + ```bash # 1. Download NCBI taxonomy taxembed-download @@ -108,6 +107,7 @@ taxembed-visualize taxonomy_model_small_best.pth ``` **Using Python scripts directly**: + ```bash python prepare_taxonomy_data.py # Download python build_transitive_closure.py # Prepare @@ -134,50 +134,57 @@ python scripts/visualize_embeddings.py my_model.pth --highlight mammals ## 📊 What's Different from Facebook's Implementation? -| Feature | Facebook | This Project | -|---------|----------|--------------| -| **Training Data** | Parent-child only | All ancestor-descendant pairs (9.8x more) | -| **Initialization** | Random | Depth-aware (root near center, leaves near boundary) | -| **Regularization** | None | Radial penalty to enforce depth → radius mapping | -| **Negative Sampling** | Random | Hard negatives (cousins at same taxonomic level) | -| **Loss Weighting** | Uniform | Depth-weighted (deeper pairs more important) | -| **Ball Constraints** | Soft projection | 3-layer enforcement (100% compliance) | -| **Performance** | Baseline | 1000x faster regularizer, 30x faster projection | +| Feature | Facebook | This Project | +| --------------------- | ----------------- | ---------------------------------------------------- | +| **Training Data** | Parent-child only | All ancestor-descendant pairs (9.8x more) | +| **Initialization** | Random | Depth-aware (root near center, leaves near boundary) | +| **Regularization** | None | Radial penalty to enforce depth → radius mapping | +| **Negative Sampling** | Random | Hard negatives (cousins at same taxonomic level) | +| **Loss Weighting** | Uniform | Depth-weighted (deeper pairs more important) | +| **Ball Constraints** | Soft projection | 3-layer enforcement (100% compliance) | +| **Performance** | Baseline | 1000x faster regularizer, 30x faster projection | --- ## 📁 Project Structure ``` -poincare-embeddings/ -├── train_hierarchical.py # Main hierarchical training script -├── build_transitive_closure.py # Generate ancestor-descendant pairs -├── analyze_hierarchy_hyperbolic.py # Evaluate hierarchy quality -├── sanity_check.py # Comprehensive validation -├── prepare_taxonomy_data.py # Download NCBI taxonomy -├── remap_edges.py # Map TaxIDs to indices +taxembed/ +├── src/taxembed/ # Main package +│ ├── models/ # Poincaré embedding models +│ ├── training/ # Training loop and data loaders +│ ├── data/ # Data downloading and processing +│ ├── visualization/ # UMAP and plotting utilities +│ ├── analysis/ # Hierarchy quality analysis +│ ├── validation/ # Sanity checks and validation +│ ├── builders/ # Dataset builders (TaxoPy) +│ └── cli/ # Command-line interface │ -├── data/ # Data files (gitignored) -│ ├── taxonomy_edges_small.edgelist -│ ├── taxonomy_edges_small_transitive.pkl -│ └── taxonomy_edges_small.mapping.tsv +├── tests/ # Comprehensive test suite +│ ├── test_models.py +│ ├── test_training.py +│ ├── test_data.py +│ └── conftest.py │ -├── scripts/ # Utility scripts -│ ├── visualize_embeddings.py -│ ├── validate_data.py -│ └── ... +├── examples/ # Example scripts +│ ├── basic_training.py +│ ├── custom_dataset.py +│ └── nn_demo.py │ -├── hype/ # Original Facebook implementation -│ ├── graph.py -│ ├── manifolds/ -│ └── ... +├── docs/ # Documentation +│ ├── user-guide.md # Comprehensive usage guide +│ ├── theory.md # Mathematical background +│ └── CLI_COMMANDS.md # CLI reference │ -├── docs/ # Documentation -│ └── archive/ # Intermediate development docs +├── _vendor/ # Facebook's original code (backup) +│ ├── hype/ +│ └── embed.py │ -├── JOURNEY.md # Development history -├── QUICKSTART.md # 5-minute guide -└── README.md # This file +├── data/ # Data files (gitignored) +├── artifacts/ # Training outputs (gitignored) +├── pyproject.toml # Package configuration +├── CONTRIBUTING.md # Contribution guidelines +└── README.md # This file ``` --- @@ -185,6 +192,7 @@ poincare-embeddings/ ## 🎯 Current Status ### **What Works ✅** + - ✅ Clean data pipeline with validation - ✅ Transitive closure computation (975K pairs) - ✅ Hierarchical training features implemented @@ -193,6 +201,7 @@ poincare-embeddings/ - ✅ Automatic checkpointing and early stopping ### **What Needs Work ⚠️** + - ⚠️ Hierarchy quality is poor after limited training (2 epochs) - ⚠️ Depth-norm correlation ~0 (should be >0.5) - ⚠️ Taxonomic separation ratios <1.1x (should be >1.5x) @@ -205,6 +214,7 @@ poincare-embeddings/ ## 🔧 Key Scripts ### **Training** + ```bash # Hierarchical training with all features python train_hierarchical.py --help @@ -214,6 +224,7 @@ python embed.py -dset data/taxonomy_edges.mapped.edgelist ... ``` ### **Analysis** + ```bash # Validate data quality python sanity_check.py @@ -226,6 +237,7 @@ python scripts/visualize_embeddings.py model.pth --highlight primates ``` ### **Data Preparation** + ```bash # Download NCBI taxonomy python prepare_taxonomy_data.py @@ -258,6 +270,7 @@ taxembed visualize echinoderms --children 1 # Color by grandchildren (--childre ``` **Features:** + - **Automatic dataset building**: Uses [TaxoPy](https://pypi.org/project/taxopy/) to query NCBI taxonomy and build datasets on-the-fly - **Smart checkpoint selection**: Visualization automatically uses the best checkpoint for each tag - **Hierarchical coloring**: `--children` flag controls depth (0=children, 1=grandchildren, 2=great-grandchildren, etc.) @@ -265,6 +278,7 @@ taxembed visualize echinoderms --children 1 # Color by grandchildren (--childre - **Organized artifacts**: All outputs stored in `artifacts/tags//` with full metadata **Advanced usage:** + ```bash # Use pre-built dataset files taxembed train --file data/my_transitive.pkl --mapping data/my.mapping.tsv -as custom_tag @@ -274,6 +288,7 @@ taxembed visualize cnidaria --sample 15000 --output custom_plot.png --root-taxid ``` ### **Build Custom Clade Datasets (Standalone)** + ```bash # Example: build the Metazoa (animals) subset with automatic mapping uv run python scripts/build_clade_dataset.py \ @@ -282,6 +297,7 @@ uv run python scripts/build_clade_dataset.py \ ``` This leverages [TaxoPy](https://pypi.org/project/taxopy/) to: + - Query NCBI taxonomy for all descendants of the requested root - Emit raw and remapped edgelists (`data/taxopy//taxonomy_edges_.edgelist`) - Write mapping + manifest files for reproducible provenance @@ -293,10 +309,11 @@ Use `--max-depth` to truncate deep subtrees or point `--taxdump-dir` at an alter ## 📖 Documentation -- **[QUICKSTART.md](QUICKSTART.md)** - Get started in 5 minutes -- **[JOURNEY.md](JOURNEY.md)** - Full development history from Facebook's code to now -- **[SESSION_SUMMARY_NOV8.md](SESSION_SUMMARY_NOV8.md)** - Latest session summary with findings -- **[docs/archive/](docs/archive/)** - Intermediate development documents +- **[docs/user-guide.md](docs/user-guide.md)** - Comprehensive usage guide +- **[docs/theory.md](docs/theory.md)** - Mathematical background and theory +- **[docs/CLI_COMMANDS.md](docs/CLI_COMMANDS.md)** - Command-line reference +- **[examples/](examples/)** - Example scripts and tutorials +- **[CONTRIBUTING.md](CONTRIBUTING.md)** - How to contribute --- @@ -309,6 +326,7 @@ python sanity_check.py ``` This validates: + - ✅ Mapping file integrity (no duplicates, continuous indices) - ✅ Transitive closure data (valid indices, no self-loops) - ✅ Projection logic (keeps embeddings in ball) @@ -325,12 +343,14 @@ This validates: ## 📈 Performance ### **Optimizations Applied** + - **Regularizer**: Vectorized (1000x faster, 1.7B → 111K ops/epoch) - **Projection**: Selective (30x faster, only updated embeddings) - **Tensor Creation**: Pre-allocated arrays (10-100x faster) - **Device**: CPU-only on macOS (stable, no MPS hanging) ### **Training Speed** + - Small dataset (111K organisms): ~3 min/epoch on M3 Mac - Full dataset (2.7M organisms): ~60 min/epoch on M3 Mac @@ -339,18 +359,20 @@ This validates: ## 🔬 Experimental Results ### **Ball Constraint Enforcement** -| Version | Max Norm | Outside Ball | Status | -|---------|----------|--------------|--------| -| v1 (weak reg) | 2.18 | 54% | ❌ Broken | -| v2 (strong reg) | 1.45 | 2.2% | ⚠️ Better | -| v3 (3-layer) | 1.00 | 0% | ✅ Perfect | + +| Version | Max Norm | Outside Ball | Status | +| --------------- | -------- | ------------ | ---------- | +| v1 (weak reg) | 2.18 | 54% | ❌ Broken | +| v2 (strong reg) | 1.45 | 2.2% | ⚠️ Better | +| v3 (3-layer) | 1.00 | 0% | ✅ Perfect | ### **Hierarchy Quality** (After 2 epochs) -| Metric | Target | Actual | Status | -|--------|--------|--------|--------| -| Depth-norm corr | >0.5 | +0.003 | ❌ Poor | -| Phylum sep | >1.5x | 1.08x | ❌ Poor | -| Class sep | >1.5x | 0.99x | ❌ Poor | + +| Metric | Target | Actual | Status | +| --------------- | ------ | ------ | ------- | +| Depth-norm corr | >0.5 | +0.003 | ❌ Poor | +| Phylum sep | >1.5x | 1.08x | ❌ Poor | +| Class sep | >1.5x | 0.99x | ❌ Poor | **Conclusion:** Constraints work perfectly, but hierarchy learning needs more time or tuning. @@ -359,12 +381,14 @@ This validates: ## 🚧 Known Issues & Future Work ### **Current Limitations** + 1. **Poor hierarchy quality** - Only 2 epochs completed, needs more training 2. **Data imbalance** - 94% deep ancestors, 6% parent-child (may need balanced sampling) 3. **Regularization trade-off** - λ=0.1 enforces constraints but may limit expressiveness 4. **No curriculum learning** - Trains on all pairs at once (may need progressive training) ### **Future Directions** + 1. Train longer with increased patience (50-100 epochs) 2. Implement balanced sampling (equal parent-child, grandparent, deep) 3. Progressive training (parent-child → grandparent → all ancestors) @@ -373,11 +397,46 @@ This validates: --- +## 🛠️ Development + +### **Code Quality Tools** + +This project maintains high code quality using modern Python tooling: + +```bash +# Linting and formatting with Ruff +uv run ruff check . # Check for issues +uv run ruff check --fix . # Auto-fix issues +uv run ruff format . # Format code + +# Static type checking with MyPy +uv run mypy src/taxembed # Type check source + +# Testing with Pytest +uv run pytest # Run test suite +uv run pytest --cov=src/taxembed # With coverage + +# Complete quality check +uv run ruff check . && uv run mypy src/taxembed && uv run pytest +``` + +**Configuration:** + +- All tools configured in `pyproject.toml` +- Ruff: 100 char lines, Python 3.11+, comprehensive rules +- MyPy: Strict typing with gradual adoption strategy +- Pytest: Comprehensive test suite with coverage reporting + +See [CONTRIBUTING.md](CONTRIBUTING.md) for detailed guidelines. + +--- + ## 🤝 Contributing Contributions are welcome! Please read [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines. ### **Priority Areas** + - Hyperparameter tuning for better hierarchy quality - Balanced/curriculum sampling strategies - Alternative hyperbolic models (Lorentz, Klein) @@ -389,14 +448,17 @@ Contributions are welcome! Please read [CONTRIBUTING.md](CONTRIBUTING.md) for gu ## 📚 References ### **Original Papers** + - Nickel & Kiela (2017). "Poincaré Embeddings for Learning Hierarchical Representations" [[PDF](https://arxiv.org/abs/1705.08039)] - Facebook Research implementation: [[GitHub](https://github.com/facebookresearch/poincare-embeddings)] ### **Data** + - NCBI Taxonomy: https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/ - Taxonomy documentation: https://www.ncbi.nlm.nih.gov/taxonomy ### **Related Work** + - Hyperbolic Neural Networks - Lorentz Embeddings - Box Embeddings for Hierarchies @@ -427,4 +489,4 @@ MIT License - see [LICENSE](LICENSE) file for details. **⭐ If you find this useful, please star the repository!** -*Last Updated: December 2025* +_Last Updated: December 2025_ diff --git a/_vendor/README.md b/_vendor/README.md new file mode 100644 index 0000000..1e0318a --- /dev/null +++ b/_vendor/README.md @@ -0,0 +1,26 @@ +# Vendor: Facebook's Original Poincaré Embeddings + +This directory contains the original implementation from Facebook Research, preserved for reference. + +## Source + +- **Paper**: "Poincaré Embeddings for Learning Hierarchical Representations" (Nickel & Kiela, 2017) +- **Original Repository**: https://github.com/facebookresearch/poincare-embeddings +- **arXiv**: https://arxiv.org/abs/1705.08039 + +## Contents + +- `hype/` - Original Facebook implementation package +- `embed.py` - Original training script + +## Note + +This code is kept as a backup and reference. The taxembed package has been completely rewritten with: + +- Hierarchical training features +- Transitive closure support +- Depth-aware initialization +- Modern Python packaging +- Clean CLI interface + +Do not modify files in this directory. They are preserved for historical reference only. diff --git a/embed.py b/_vendor/embed.py similarity index 100% rename from embed.py rename to _vendor/embed.py diff --git a/hype/__init__.py b/_vendor/hype/__init__.py similarity index 100% rename from hype/__init__.py rename to _vendor/hype/__init__.py diff --git a/hype/adjacency_matrix_dataset.pyi b/_vendor/hype/adjacency_matrix_dataset.pyi similarity index 100% rename from hype/adjacency_matrix_dataset.pyi rename to _vendor/hype/adjacency_matrix_dataset.pyi diff --git a/hype/adjacency_matrix_dataset.pyx b/_vendor/hype/adjacency_matrix_dataset.pyx similarity index 100% rename from hype/adjacency_matrix_dataset.pyx rename to _vendor/hype/adjacency_matrix_dataset.pyx diff --git a/hype/checkpoint.py b/_vendor/hype/checkpoint.py similarity index 100% rename from hype/checkpoint.py rename to _vendor/hype/checkpoint.py diff --git a/hype/common.py b/_vendor/hype/common.py similarity index 100% rename from hype/common.py rename to _vendor/hype/common.py diff --git a/hype/energy_function.py b/_vendor/hype/energy_function.py similarity index 100% rename from hype/energy_function.py rename to _vendor/hype/energy_function.py diff --git a/hype/graph.py b/_vendor/hype/graph.py similarity index 100% rename from hype/graph.py rename to _vendor/hype/graph.py diff --git a/hype/graph_dataset.pyx b/_vendor/hype/graph_dataset.pyx similarity index 100% rename from hype/graph_dataset.pyx rename to _vendor/hype/graph_dataset.pyx diff --git a/hype/hypernymy_eval.py b/_vendor/hype/hypernymy_eval.py similarity index 100% rename from hype/hypernymy_eval.py rename to _vendor/hype/hypernymy_eval.py diff --git a/hype/manifolds/__init__.py b/_vendor/hype/manifolds/__init__.py similarity index 100% rename from hype/manifolds/__init__.py rename to _vendor/hype/manifolds/__init__.py diff --git a/hype/manifolds/euclidean.py b/_vendor/hype/manifolds/euclidean.py similarity index 100% rename from hype/manifolds/euclidean.py rename to _vendor/hype/manifolds/euclidean.py diff --git a/hype/manifolds/lorentz.py b/_vendor/hype/manifolds/lorentz.py similarity index 100% rename from hype/manifolds/lorentz.py rename to _vendor/hype/manifolds/lorentz.py diff --git a/hype/manifolds/manifold.py b/_vendor/hype/manifolds/manifold.py similarity index 100% rename from hype/manifolds/manifold.py rename to _vendor/hype/manifolds/manifold.py diff --git a/hype/manifolds/poincare.py b/_vendor/hype/manifolds/poincare.py similarity index 100% rename from hype/manifolds/poincare.py rename to _vendor/hype/manifolds/poincare.py diff --git a/hype/path_manager.py b/_vendor/hype/path_manager.py similarity index 100% rename from hype/path_manager.py rename to _vendor/hype/path_manager.py diff --git a/hype/rsgd.py b/_vendor/hype/rsgd.py similarity index 100% rename from hype/rsgd.py rename to _vendor/hype/rsgd.py diff --git a/hype/sn.py b/_vendor/hype/sn.py similarity index 100% rename from hype/sn.py rename to _vendor/hype/sn.py diff --git a/hype/train.py b/_vendor/hype/train.py similarity index 100% rename from hype/train.py rename to _vendor/hype/train.py diff --git a/analyze_hierarchy.py b/analyze_hierarchy.py deleted file mode 100644 index f38be7b..0000000 --- a/analyze_hierarchy.py +++ /dev/null @@ -1,312 +0,0 @@ -#!/usr/bin/env python3 -""" -Analyze hierarchical structure in embeddings. -Check if organisms within the same high-level taxon (phylum/class/order) -are closer to each other than to organisms from different taxa. -""" - -import torch -import numpy as np -import pandas as pd -from collections import defaultdict -from scipy.spatial.distance import pdist, squareform -import matplotlib.pyplot as plt - - -def load_embeddings(ckpt_path): - """Load embeddings from checkpoint.""" - print(f"Loading embeddings from {ckpt_path}...") - ckpt = torch.load(ckpt_path, map_location="cpu") - - if "state_dict" in ckpt: - sd = ckpt["state_dict"] - emb = sd["lt.weight"].detach().cpu().numpy() - elif "embeddings" in ckpt: - emb = ckpt["embeddings"].cpu().numpy() - else: - raise ValueError("Cannot find embeddings in checkpoint") - - print(f" ✓ Shape: {emb.shape}") - return emb - - -def load_mapping(mapping_file): - """Load index to TaxID mapping.""" - print(f"Loading mapping from {mapping_file}...") - df = pd.read_csv(mapping_file, sep="\t", header=None, names=["idx", "taxid"]) - numeric_df = df[df["taxid"].str.isnumeric()] - idx2tax = dict(zip(numeric_df["idx"], numeric_df["taxid"])) - print(f" ✓ Loaded {len(idx2tax):,} mappings") - return idx2tax - - -def load_taxonomy_ranks(valid_taxids): - """Load taxonomy with ranks for valid TaxIDs.""" - print("Loading taxonomy tree with ranks...") - - # Load names - names = {} - with open("data/names.dmp", "r") as f: - for line in f: - parts = [p.strip() for p in line.split("|")] - if len(parts) >= 4 and parts[3] == "scientific name": - taxid = int(parts[0]) - if taxid in valid_taxids: - names[taxid] = parts[1] - - # Load nodes with ranks - taxonomy = {} - with open("data/nodes.dmp", "r") as f: - for line in f: - parts = [p.strip() for p in line.split("|")] - if len(parts) >= 5: - taxid = int(parts[0]) - if taxid in valid_taxids: - parent = int(parts[1]) - rank = parts[2] - taxonomy[taxid] = { - "parent": parent, - "rank": rank, - "name": names.get(taxid, f"TaxID_{taxid}") - } - - print(f" ✓ Loaded {len(taxonomy):,} taxonomy nodes") - return taxonomy - - -def get_ancestor_at_rank(taxid, taxonomy, target_rank): - """Find the ancestor of a taxid at a specific rank.""" - visited = set() - current = taxid - - while current in taxonomy and current not in visited: - visited.add(current) - node = taxonomy[current] - - if node["rank"] == target_rank: - return current - - parent = node["parent"] - if parent == current: # Root - break - current = parent - - return None - - -def analyze_hierarchical_clustering(emb, idx2tax, taxonomy, rank="phylum"): - """Analyze if organisms cluster by taxonomic rank.""" - print(f"\n{'='*80}") - print(f"HIERARCHICAL CLUSTERING ANALYSIS - {rank.upper()}") - print(f"{'='*80}\n") - - # Get valid taxids - valid_taxids = set(int(t) for t in idx2tax.values()) - - # Map each organism to its ancestor at target rank - organism_to_group = {} - group_names = {} - - for idx, taxid_str in idx2tax.items(): - taxid = int(taxid_str) - ancestor = get_ancestor_at_rank(taxid, taxonomy, rank) - if ancestor: - organism_to_group[idx] = ancestor - if ancestor not in group_names and ancestor in taxonomy: - group_names[ancestor] = taxonomy[ancestor]["name"] - - print(f"Found {len(set(organism_to_group.values()))} distinct {rank}s") - print(f"Mapped {len(organism_to_group)}/{len(idx2tax)} organisms to {rank} level") - - # Group organisms by their ancestor - # Only include indices that are valid for the embedding matrix - max_idx = emb.shape[0] - 1 - groups = defaultdict(list) - for idx, group_id in organism_to_group.items(): - idx_int = int(idx) if isinstance(idx, str) else idx - if idx_int <= max_idx: - groups[group_id].append(idx_int) - else: - print(f"Warning: idx {idx} out of bounds (max={max_idx})") - - # Filter to groups with at least 10 members for meaningful analysis - min_size = 10 - large_groups = {gid: indices for gid, indices in groups.items() if len(indices) >= min_size} - - print(f"\n{rank.capitalize()}s with ≥{min_size} organisms: {len(large_groups)}") - - # Show top groups by size - group_sizes = [(gid, len(indices), group_names.get(gid, f"TaxID_{gid}")) - for gid, indices in large_groups.items()] - group_sizes.sort(key=lambda x: -x[1]) - - print(f"\nTop 20 {rank}s by organism count:") - for i, (gid, size, name) in enumerate(group_sizes[:20], 1): - print(f" {i:2d}. {name:40s}: {size:6,} organisms") - - # Compute distances - print(f"\nComputing pairwise distances...") - - # Sample for efficiency if needed - max_per_group = 500 - sampled_groups = {} - for gid, indices in large_groups.items(): - indices_list = list(indices) if not isinstance(indices, list) else indices - if len(indices_list) > max_per_group: - sampled_groups[gid] = list(np.random.choice(indices_list, max_per_group, replace=False)) - else: - sampled_groups[gid] = indices_list - - # Calculate intra-group vs inter-group distances - intra_distances = [] - inter_distances = [] - - group_list = list(sampled_groups.items()) - - for i, (gid1, indices1) in enumerate(group_list): - # Intra-group distances - if len(indices1) >= 2: - # Convert to numpy array of embeddings - group_emb = np.array([emb[int(idx)] for idx in indices1]) - dists = pdist(group_emb) - intra_distances.extend(dists) - - # Inter-group distances (sample to avoid O(n²) explosion) - for j, (gid2, indices2) in enumerate(group_list[i+1:], i+1): - # Sample pairs for efficiency - idx1_sample = list(np.random.choice(indices1, min(100, len(indices1)), replace=False)) - idx2_sample = list(np.random.choice(indices2, min(100, len(indices2)), replace=False)) - - for idx1 in idx1_sample[:10]: # Limit per group pair - for idx2 in idx2_sample[:10]: - dist = np.linalg.norm(emb[int(idx1)] - emb[int(idx2)]) - inter_distances.append(dist) - - intra_distances = np.array(intra_distances) - inter_distances = np.array(inter_distances) - - print(f"\n{'='*80}") - print(f"DISTANCE STATISTICS") - print(f"{'='*80}") - print(f"\nIntra-{rank} distances (within same {rank}):") - print(f" Count: {len(intra_distances):,}") - print(f" Mean: {np.mean(intra_distances):.6f}") - print(f" Std: {np.std(intra_distances):.6f}") - print(f" Min: {np.min(intra_distances):.6f}") - print(f" Max: {np.max(intra_distances):.6f}") - print(f" Median: {np.median(intra_distances):.6f}") - - print(f"\nInter-{rank} distances (between different {rank}s):") - print(f" Count: {len(inter_distances):,}") - print(f" Mean: {np.mean(inter_distances):.6f}") - print(f" Std: {np.std(inter_distances):.6f}") - print(f" Min: {np.min(inter_distances):.6f}") - print(f" Max: {np.max(inter_distances):.6f}") - print(f" Median: {np.median(inter_distances):.6f}") - - # Separation score - separation = np.mean(inter_distances) / np.mean(intra_distances) - print(f"\n{'='*80}") - print(f"HIERARCHICAL QUALITY METRICS") - print(f"{'='*80}") - print(f"\nSeparation Ratio (higher is better):") - print(f" Inter-{rank} / Intra-{rank} = {separation:.3f}x") - - if separation > 2.0: - quality = "✅ EXCELLENT" - elif separation > 1.5: - quality = "✅ GOOD" - elif separation > 1.2: - quality = "⚠️ MODERATE" - else: - quality = "❌ POOR" - - print(f" Quality: {quality}") - - # Overlap analysis - overlap = np.sum(intra_distances > np.median(inter_distances)) / len(intra_distances) * 100 - print(f"\nDistance overlap:") - print(f" {overlap:.1f}% of intra-{rank} distances exceed median inter-{rank} distance") - - # Plot - fig, axes = plt.subplots(1, 2, figsize=(16, 6)) - - # Histogram - ax = axes[0] - bins = np.linspace(0, max(np.max(intra_distances), np.max(inter_distances)), 50) - ax.hist(intra_distances, bins=bins, alpha=0.6, label=f'Intra-{rank}', color='blue', density=True) - ax.hist(inter_distances, bins=bins, alpha=0.6, label=f'Inter-{rank}', color='red', density=True) - ax.axvline(np.mean(intra_distances), color='blue', linestyle='--', linewidth=2, label=f'Intra mean: {np.mean(intra_distances):.3f}') - ax.axvline(np.mean(inter_distances), color='red', linestyle='--', linewidth=2, label=f'Inter mean: {np.mean(inter_distances):.3f}') - ax.set_xlabel('Euclidean Distance', fontsize=12) - ax.set_ylabel('Density', fontsize=12) - ax.set_title(f'Distance Distribution by {rank.capitalize()}', fontsize=14, fontweight='bold') - ax.legend() - ax.grid(True, alpha=0.3) - - # Box plot - ax = axes[1] - data_to_plot = [intra_distances, inter_distances] - bp = ax.boxplot(data_to_plot, labels=[f'Intra-{rank}', f'Inter-{rank}'], - patch_artist=True, widths=0.6) - bp['boxes'][0].set_facecolor('lightblue') - bp['boxes'][1].set_facecolor('lightcoral') - ax.set_ylabel('Euclidean Distance', fontsize=12) - ax.set_title(f'{rank.capitalize()}-level Clustering Quality', fontsize=14, fontweight='bold') - ax.grid(True, alpha=0.3, axis='y') - - # Add separation ratio text - ax.text(0.5, 0.95, f'Separation Ratio: {separation:.2f}x\n{quality}', - transform=ax.transAxes, fontsize=12, verticalalignment='top', - bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5), - horizontalalignment='center') - - plt.tight_layout() - output_file = f'hierarchy_analysis_{rank}.png' - plt.savefig(output_file, dpi=150, bbox_inches='tight') - print(f"\nSaved plot: {output_file}") - plt.close() - - return { - 'intra_mean': np.mean(intra_distances), - 'inter_mean': np.mean(inter_distances), - 'separation': separation, - 'quality': quality, - 'num_groups': len(large_groups), - 'num_organisms': sum(len(indices) for indices in large_groups.values()) - } - - -def main(): - # Load data - checkpoint = "taxonomy_model_small_early_stop_epoch2341.pth" - mapping_file = "data/taxonomy_edges_small.mapping.tsv" - - emb = load_embeddings(checkpoint) - idx2tax = load_mapping(mapping_file) - - # Load taxonomy - valid_taxids = set(int(t) for t in idx2tax.values()) - taxonomy = load_taxonomy_ranks(valid_taxids) - - # Analyze at different taxonomic levels - results = {} - for rank in ["phylum", "class", "order", "family"]: - try: - results[rank] = analyze_hierarchical_clustering(emb, idx2tax, taxonomy, rank=rank) - except Exception as e: - print(f"\nError analyzing {rank}: {e}") - - # Summary - print(f"\n{'='*80}") - print(f"SUMMARY ACROSS RANKS") - print(f"{'='*80}\n") - print(f"{'Rank':<12} {'Groups':>8} {'Organisms':>10} {'Separation':>12} {'Quality':>15}") - print("-" * 80) - for rank, res in results.items(): - print(f"{rank.capitalize():<12} {res['num_groups']:>8} {res['num_organisms']:>10,} " - f"{res['separation']:>11.2f}x {res['quality']:>15}") - - -if __name__ == "__main__": - main() diff --git a/check_dataset_composition.py b/check_dataset_composition.py deleted file mode 100644 index be81e97..0000000 --- a/check_dataset_composition.py +++ /dev/null @@ -1,101 +0,0 @@ -#!/usr/bin/env python3 -""" -Check the composition of organisms in the dataset by taxonomic group. -""" - -import pandas as pd -from collections import defaultdict - -def load_taxonomy_tree(valid_taxids): - """Load taxonomy tree filtered to valid TaxIDs.""" - names = {} - with open("data/names.dmp", "r") as f: - for line in f: - parts = [p.strip() for p in line.split("|")] - if len(parts) >= 4 and parts[3] == "scientific name": - taxid = int(parts[0]) - if taxid in valid_taxids: - names[taxid] = parts[1] - - nodes = {} - with open("data/nodes.dmp", "r") as f: - for line in f: - parts = [p.strip() for p in line.split("|")] - if len(parts) >= 5: - taxid = int(parts[0]) - if taxid in valid_taxids: - parent = int(parts[1]) - rank = parts[2] - nodes[taxid] = {"parent": parent, "rank": rank, "name": names.get(taxid, "")} - - return nodes - - -def find_group_descendants(nodes, root_taxid): - """Find all descendants of a root taxid.""" - descendants = set() - - def find_desc(taxid): - descendants.add(taxid) - for child_id, child_info in nodes.items(): - if child_info["parent"] == taxid and child_id not in descendants: - find_desc(child_id) - - if root_taxid in nodes: - find_desc(root_taxid) - - return descendants - - -def main(): - # Load mapping - df = pd.read_csv("data/taxonomy_edges_small.mapping.tsv", sep="\t", header=None, names=["idx", "taxid"]) - valid_taxids = set(int(x) for x in df["taxid"] if str(x).isnumeric()) - - print(f"Total organisms in small dataset: {len(valid_taxids):,}\n") - - # Load taxonomy - nodes = load_taxonomy_tree(valid_taxids) - - # Check major taxonomic groups - groups = { - "Primates": 9443, - "Mammals": 40674, - "Vertebrates": 7742, - "Bacteria": 2, - "Archaea": 2157, - "Fungi": 4751, - "Plants (Viridiplantae)": 33090, - "Insects": 50557, - "Rodents": 9989, - "Nematodes": 6231, - "Arthropods": 6656, - "Metazoa (Animals)": 33208, - } - - print("=" * 70) - print("DATASET COMPOSITION BY TAXONOMIC GROUP") - print("=" * 70) - - for group_name, root_taxid in groups.items(): - descendants = find_group_descendants(nodes, root_taxid) - count = len(descendants) - percentage = (count / len(valid_taxids)) * 100 - print(f"{group_name:30s}: {count:6,} organisms ({percentage:5.2f}%)") - - print("=" * 70) - - # Check top-level domains - print("\nTop-level taxonomy distribution:") - rank_counts = defaultdict(int) - for taxid, info in nodes.items(): - rank = info.get("rank", "unknown") - rank_counts[rank] += 1 - - for rank, count in sorted(rank_counts.items(), key=lambda x: -x[1])[:15]: - percentage = (count / len(valid_taxids)) * 100 - print(f" {rank:20s}: {count:6,} ({percentage:5.2f}%)") - - -if __name__ == "__main__": - main() diff --git a/check_model.py b/check_model.py deleted file mode 100644 index 4a03ec3..0000000 --- a/check_model.py +++ /dev/null @@ -1,91 +0,0 @@ -#!/usr/bin/env python3 -"""Quick script to check trained model status.""" - -import torch -import os -from glob import glob - -print("="*80) -print("TRAINED MODEL SUMMARY") -print("="*80) - -# Find all model files -model_files = sorted(glob("taxonomy_model_small*.pth")) - -if not model_files: - print("❌ No trained models found!") - exit(1) - -print(f"\nFound {len(model_files)} model files:\n") - -best_loss = float('inf') -best_file = None - -for model_file in model_files: - size_mb = os.path.getsize(model_file) / (1024**2) - - # Load checkpoint - try: - ckpt = torch.load(model_file, map_location='cpu') - - # Extract info - epoch = ckpt.get('epoch', '?') - loss = ckpt.get('loss', None) - embeddings = ckpt.get('embeddings', None) - - if embeddings is not None: - n_nodes = embeddings.shape[0] - dim = embeddings.shape[1] - norms = embeddings.norm(dim=1) - max_norm = norms.max().item() - mean_norm = norms.mean().item() - outside = (norms >= 1.0).sum().item() - else: - n_nodes = dim = max_norm = mean_norm = outside = "?" - - # Track best - if loss is not None and loss < best_loss: - best_loss = loss - best_file = model_file - - # Print info - status = "✅" if outside == 0 else f"⚠️ {outside} outside" - loss_str = f"{loss:.6f}" if loss is not None else "N/A" - - print(f" {os.path.basename(model_file):35s} | " - f"Epoch {epoch:>3} | " - f"Loss: {loss_str:>10} | " - f"Nodes: {n_nodes:>6,} | " - f"Dim: {dim:>2} | " - f"MaxNorm: {max_norm:.4f} | " - f"{status}") - - except Exception as e: - print(f" {os.path.basename(model_file):35s} | ERROR: {e}") - -print("\n" + "="*80) -if best_file: - print(f"🏆 BEST MODEL: {best_file}") - print(f" Loss: {best_loss:.6f}") - - # Create symlink or copy - best_link = "taxonomy_model_small_best.pth" - if os.path.exists(best_link): - os.remove(best_link) - - # Create a copy - import shutil - shutil.copy(best_file, best_link) - print(f" ✓ Copied to: {best_link}") -else: - print("⚠️ Could not determine best model") - -print("="*80) -print("\nRECOMMENDED NEXT STEPS:") -print(" 1. Analyze hierarchy quality:") -print(" python analyze_hierarchy_hyperbolic.py") -print("\n 2. Visualize embeddings:") -print(" python scripts/visualize_embeddings.py taxonomy_model_small_best.pth --highlight mammals") -print("\n 3. Query nearest neighbors:") -print(" python query_embeddings.py taxonomy_model_small_best.pth") -print("="*80) diff --git a/docs/CLEANUP_SUMMARY.md b/docs/CLEANUP_SUMMARY.md deleted file mode 100644 index 5e35efa..0000000 --- a/docs/CLEANUP_SUMMARY.md +++ /dev/null @@ -1,281 +0,0 @@ -# Repository Cleanup Summary - -**Date:** November 12, 2025 - -## ✅ Cleanup Complete - -The repository has been modernized and organized with proper Python packaging standards. - ---- - -## 🗑️ Removed Files - -### Legacy Facebook Research Files -- `wn-nouns.jpg` - WordNet visualization -- `README.org` - Original Emacs org-mode readme -- `wordnet/` directory - WordNet-specific scripts -- `hypernymy_eval.py` - WordNet evaluation -- `reconstruction.py` - WordNet reconstruction -- `environment.yml` - Conda environment (replaced by uv) - -### Old Build System -- `setup.py` - Replaced by modern `pyproject.toml` with hatchling - -### Redundant Scripts -- `cleanup_repo.sh` - Old cleanup script -- `cleanup_for_release.sh` - Release cleanup -- `cleanup_old_checkpoints.py` - Checkpoint cleanup -- `git_push_commands.sh` - Git automation -- `watch_training.sh` - Training monitor -- `run_hierarchical_training.sh` - Old training wrapper - -### Duplicate Analysis Scripts -- `assess_training.py` - Redundant with `check_model.py` -- `monitor_training.py` - Superseded by `train_small.py` built-in metrics -- `resume_training.py` - Functionality in main training scripts -- `train_with_early_stopping.py` - Merged into `train_small.py` -- `evaluate_full.py` - Redundant -- `evaluate_and_visualize.py` - Split into focused scripts - -### Duplicate Model Files in Root -- `taxonomy_model_small.pth` - Kept in `small_model_28epoch/` -- `taxonomy_model_small_best.pth` - Kept in `small_model_28epoch/` -- `taxonomy_model_small_epoch*.pth` (5 files) - Kept in `small_model_28epoch/` -- `taxonomy_embeddings_multi_groups.png` - Kept in `small_model_28epoch/` - ---- - -## 📁 New Structure - -### Root Directory (Clean!) -``` -poincare-embeddings/ -├── README.md # Main documentation -├── QUICKSTART.md # Quick start guide -├── LICENSE # MIT license -├── pyproject.toml # Modern Python packaging (hatchling + uv) -├── ruff.toml # Code quality config -├── Makefile # Common commands -└── requirements.txt # Fallback pip requirements -``` - -### Documentation (docs/) -All supplementary documentation moved here: -``` -docs/ -├── JOURNEY.md # Development history (8 phases) -├── FINAL_STATUS.md # Production status -├── TRAIN_SMALL_GUIDE.md # Training guide -├── TRAIN_FULL_GUIDE.md # Full dataset reference -├── COMMIT_SUMMARY.md # Commit information -├── RELEASE_SUMMARY.md # Release notes -├── CONTRIBUTING.md # Contribution guidelines -├── CODE_OF_CONDUCT.md # Community standards -├── PRE_PUSH_CHECKLIST.md # Pre-push checklist -└── archive/ # Historical documents -``` - -### Core Scripts (Root) -Focused, essential scripts: -``` -├── train_small.py # Main training script ⭐ -├── train_hierarchical.py # Core hierarchical model -├── visualize_multi_groups.py # UMAP visualization -├── build_transitive_closure.py # Data preparation -├── prepare_taxonomy_data.py # NCBI download -├── remap_edges.py # ID remapping -├── check_model.py # Model analysis -├── analyze_hierarchy.py # Hierarchy analysis -├── analyze_hierarchy_hyperbolic.py # Hyperbolic analysis -├── check_dataset_composition.py # Data validation -├── final_sanity_check.py # Sanity checks -└── embed.py # Original Poincaré training -``` - -### Production Model -``` -small_model_28epoch/ -├── taxonomy_model_small_best.pth # Best model (epoch 28, loss 0.472) -├── taxonomy_embeddings_multi_groups.png -├── best_epoch_analysis_epoch28.png -└── umap_taxonomy_model_small_best_mammals_highlighted.png -``` - -### Reference Model -``` -taxonomy_model_animals_best.pth # 1M organisms (incomplete, 4 epochs) -``` - ---- - -## 🔧 Modernized Configuration - -### pyproject.toml (NEW) -- **Build system:** `hatchling` (lightweight, modern) -- **Package manager:** `uv` (10-100x faster than pip) -- **Python version:** >=3.11 -- **Dependencies:** Streamlined (torch, numpy, pandas, matplotlib, umap) -- **Dev tools:** ruff, pytest, mypy -- **Proper package:** `src/taxembed/` - -### Key Improvements: -```toml -[build-system] -requires = ["hatchling"] # Was: setuptools + cython - -[project] -requires-python = ">=3.11" # Was: >=3.8 -dependencies = [ - "torch>=2.0.0", - # Core dependencies only -] - -[tool.uv] -dev-dependencies = [ - "ruff>=0.6.0", # Latest - "pytest>=8.0.0", - "mypy>=1.0.0", -] - -[tool.ruff] -target-version = "py311" # Was: py38 -exclude = ["hype"] # Ignore original code -``` - -### Makefile (UPDATED) -```makefile -# New commands -make install # uv sync -make install-dev # uv sync --all-extras -make train # Quick test (1 epoch) -make check # Sanity checks -make lint # ruff check -make format # ruff format -make test # pytest -make clean # Remove artifacts -``` - ---- - -## 📊 Statistics - -### Files Removed: 23 -- Legacy: 6 -- Redundant scripts: 11 -- Duplicate models: 6 - -### Disk Space Freed: ~185 MB -- Duplicate checkpoints: ~160 MB -- Legacy files: ~25 MB - -### Lines of Configuration: ~100 -- Modern `pyproject.toml`: 98 lines -- Clean `Makefile`: 66 lines - ---- - -## ✨ Benefits - -### 1. Cleaner Repository -- Root has only essential files -- Clear separation: code vs docs vs data -- No legacy cruft from original repo - -### 2. Modern Python Packaging -- Standard `pyproject.toml` (PEP 621) -- Fast dependency management with `uv` -- No compilation required (removed Cython) -- Proper package structure (`src/taxembed/`) - -### 3. Better Code Quality -- `ruff` for linting and formatting -- `mypy` for type checking -- `pytest` for testing -- All configured in `pyproject.toml` - -### 4. Improved Developer Experience -- Simple `make` commands -- Clear documentation structure -- Easy onboarding (QUICKSTART.md) -- Fast installs with `uv` - -### 5. Production Ready -- Clean, professional structure -- Comprehensive documentation -- Validated and tested -- Ready for deployment - ---- - -## 🚀 Next Steps - -### For Users: -```bash -make install # Install dependencies -python train_small.py # Train model -make check # Verify installation -``` - -### For Developers: -```bash -make install-dev # Install with dev tools -make lint # Check code quality -make format # Format code -make test # Run tests -``` - -### For Contributors: -```bash -# See docs/CONTRIBUTING.md -``` - ---- - -## 📝 Migration Notes - -### If you had custom scripts: -- Check if functionality exists in new structure -- See `docs/` for equivalent commands -- Old scripts may be in `docs/archive/` - -### If you used old commands: -| Old | New | -|-----|-----| -| `bash run_hierarchical_training.sh` | `python train_small.py` | -| `python setup.py build_ext --inplace` | Not needed | -| `pip install -e .` | `make install` or `uv sync` | -| `python scripts/train.py` | `python train_small.py` | - -### If you need old files: -- Check `docs/archive/` for historical documents -- Git history preserves all removed files -- Contact maintainers if something is missing - ---- - -## ✅ Validation - -Ran comprehensive checks: -```bash -✅ final_sanity_check.py - All checks passed -✅ Core scripts present and working -✅ Documentation complete -✅ Production model validated -✅ No legacy files remaining -``` - ---- - -## 🎯 Result - -**Professional, modern, production-ready repository** with: -- ✅ Clean root directory -- ✅ Modern Python packaging (hatchling + uv) -- ✅ Code quality tools (ruff) -- ✅ Clear documentation structure -- ✅ Fast dependency management -- ✅ Ready for contribution and deployment - ---- - -*Repository cleaned and modernized on November 12, 2025* diff --git a/docs/COMMIT_SUMMARY.md b/docs/COMMIT_SUMMARY.md deleted file mode 100644 index babe6c5..0000000 --- a/docs/COMMIT_SUMMARY.md +++ /dev/null @@ -1,191 +0,0 @@ -# Commit Summary: Hierarchical Poincaré Embeddings Complete - -**Date:** November 10, 2025 - ---- - -## 🎯 Summary - -Successfully developed and validated hierarchical Poincaré embeddings for NCBI taxonomy. **Production-ready model available** for 92K organisms with excellent hierarchical structure. - ---- - -## ✅ What's Included - -### **1. Production Model** -- **Location:** `small_model_28epoch/` -- **Best epoch:** 28 -- **Loss:** 0.472 (51.6% improvement) -- **Quality:** 100% ball constraint compliance -- **Size:** 3.5 MB -- **Organisms:** 92,290 embedded in 10 dimensions - -### **2. Complete Training Pipeline** -- `train_small.py` - Main training script with fixed early stopping -- `train_hierarchical.py` - Core hierarchical model implementation -- `visualize_multi_groups.py` - Multi-group UMAP visualization -- `build_transitive_closure.py` - Transitive closure computation - -### **3. Comprehensive Documentation** -- `README.md` - Project overview -- `JOURNEY.md` - Complete development history (8 phases) -- `FINAL_STATUS.md` - Production status and usage guide -- `TRAIN_SMALL_GUIDE.md` - Training instructions - -### **4. Reference Model** -- `taxonomy_model_animals_best.pth` - 1M organisms (4 epochs, incomplete) -- Proof of scalability for future work - ---- - -## 🔧 Key Fixes Applied - -### **1. Early Stopping Bug (Critical)** -```python -# Before (WRONG): Compared against updated value -tracker.update(epoch, metrics) -if avg_loss < tracker.best_loss: # Always comparing self! - -# After (CORRECT): Save previous best first -prev_best_loss = tracker.best_loss -tracker.update(epoch, metrics) -if avg_loss < prev_best_loss: -``` -**Impact:** Allowed training to reach epoch 28 (was stopping at 5) - -### **2. Hyperbolic Geometry (Critical)** -```python -# Before (WRONG): Euclidean distance -umap.UMAP(metric='euclidean') # Treats hyperbolic as flat - -# After (CORRECT): Poincaré distance -d = arcosh(1 + 2·||x-y||²/((1-||x||²)(1-||y||²))) -umap.UMAP(metric='precomputed') -``` -**Impact:** Correct hierarchical structure visualization - -### **3. Data Quality** -- Fixed TaxID header contamination -- Corrected mapping file inconsistencies -- Added comprehensive validation - ---- - -## 📊 Results - -### **Small Model (Production)** -| Metric | Value | -|--------|-------| -| Organisms | 92,290 | -| Best Epoch | 28 | -| Loss | 0.472 | -| Improvement | 51.6% | -| Ball Constraint | 100% ✅ | -| Training Time | 2.5 hours (M3 Mac CPU) | - -### **Animals Model (Reference)** -| Metric | Value | -|--------|-------| -| Organisms | 1,055,469 | -| Epochs | 4 (incomplete) | -| Loss | 0.635 | -| Status | Proof of scalability | - ---- - -## 🎓 Key Insights - -1. **Convergence requires patience** - 28 epochs needed (not 2-5) -2. **Early stopping is dangerous** - Must implement correctly -3. **Geometry matters** - Hyperbolic embeddings need hyperbolic distance -4. **Hard negatives don't scale** - O(n²) fails beyond ~100K nodes -5. **Small datasets work best** - 111K is sweet spot for CPU training - ---- - -## 🗑️ Cleaned Up - -Removed: -- ✅ Animals model intermediate epochs (4 files, 160 MB) -- ✅ Temporary visualizations (3 files) -- ✅ Failed attempt scripts (8 files) -- ✅ Analysis temp scripts (2 files) - -Preserved: -- ✅ `small_model_28epoch/` (production model + viz) -- ✅ `taxonomy_model_animals_best.pth` (reference) -- ✅ Core training pipeline -- ✅ All documentation - ---- - -## 📝 Files Changed - -### **New Files** -- `JOURNEY.md` - Updated with phases 6-8 -- `FINAL_STATUS.md` - Complete project status -- `COMMIT_SUMMARY.md` - This file -- `cleanup_repo.sh` - Cleanup script -- `final_sanity_check.py` - Validation script - -### **Updated Files** -- `train_small.py` - Fixed early stopping bug (line 246, 274) - -### **Organized** -- `small_model_28epoch/` - All production files consolidated - ---- - -## ✅ Sanity Check Results - -All checks passed: -- ✅ Core scripts present -- ✅ Documentation complete -- ✅ Small model valid (92K organisms, loss 0.472, 100% in ball) -- ✅ Animals model valid (1M organisms, loss 0.635, 100% in ball) -- ✅ Data files intact -- ✅ No intermediate files remaining - ---- - -## 🚀 Ready For - -- ✅ Downstream ML tasks -- ✅ Taxonomic prediction -- ✅ Hierarchical queries -- ✅ Nearest neighbor search -- ✅ Transfer learning - ---- - -## 📦 Repository Structure - -``` -poincare-embeddings/ -├── small_model_28epoch/ # ⭐ Production model -│ ├── taxonomy_model_small_best.pth -│ ├── taxonomy_embeddings_multi_groups.png -│ ├── best_epoch_analysis_epoch28.png -│ └── umap_taxonomy_model_small_best_mammals_highlighted.png -├── train_small.py # ⭐ Main training script -├── train_hierarchical.py # Core model -├── visualize_multi_groups.py # Visualization -├── build_transitive_closure.py # Data prep -├── README.md # ⭐ Main docs -├── JOURNEY.md # ⭐ Development history -├── FINAL_STATUS.md # ⭐ Status & usage -├── taxonomy_model_animals_best.pth # Reference (1M organisms) -└── data/ # NCBI taxonomy -``` - ---- - -## 🏆 Status - -**✅ PRODUCTION READY** - -The small dataset model is fully validated, documented, and ready for deployment. - ---- - -*Last commit: November 10, 2025* diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md deleted file mode 100644 index d45f7fb..0000000 --- a/docs/CONTRIBUTING.md +++ /dev/null @@ -1,83 +0,0 @@ -# Contributing to taxembed - -We want to make contributing to this project as easy and transparent as possible. - -## Development Setup - -1. Clone the repository: -```bash -git clone https://github.com/jcoludar/taxembed.git -cd taxembed -``` - -2. Install development dependencies using uv: -```bash -uv sync -``` - -3. Build C++ extensions: -```bash -uv run python setup.py build_ext --inplace -``` - -## Pull Requests - -We actively welcome your pull requests. - -1. Fork the repo and create your branch from `main`. -2. If you've added code that should be tested, add tests in the `tests/` directory. -3. If you've changed APIs, update the documentation. -4. Ensure the test suite passes: `uv run pytest` -5. Make sure your code passes linting: `uv run ruff check src/ scripts/` -6. Format your code: `uv run ruff format src/ scripts/` - -## Code Quality - -This project uses **ruff** for linting and code formatting. - -### Linting - -Check for linting issues: -```bash -uv run ruff check src/ scripts/ -``` - -Fix linting issues automatically: -```bash -uv run ruff check --fix src/ scripts/ -``` - -### Formatting - -Format code to match project style: -```bash -uv run ruff format src/ scripts/ -``` - -### Testing - -Run the test suite: -```bash -uv run pytest -``` - -With coverage report: -```bash -uv run pytest --cov=src/taxembed -``` - -## Coding Style - -- Follow PEP 8 guidelines -- Use type hints where possible -- Write docstrings for all public functions and classes -- Keep lines under 100 characters (enforced by ruff) -- Use meaningful variable and function names - -## Issues - -We use GitHub issues to track public bugs. Please ensure your description is clear and has sufficient instructions to reproduce the issue. - -## License - -By contributing to taxembed, you agree that your contributions will be licensed under the CC-BY-NC 4.0 license found in the LICENSE file in the root directory of this source tree. diff --git a/docs/FINAL_STATUS.md b/docs/FINAL_STATUS.md deleted file mode 100644 index e88d643..0000000 --- a/docs/FINAL_STATUS.md +++ /dev/null @@ -1,348 +0,0 @@ -# Final Project Status - -**Last Updated:** November 10, 2025 - ---- - -## ✅ Project Complete - -Successfully developed hierarchical Poincaré embeddings for NCBI biological taxonomy with production-ready results. - ---- - -## 📦 Deliverables - -### **1. Production Model: Small Dataset (111K organisms)** - -**Location:** `small_model_28epoch/` - -**Files:** -- `taxonomy_model_small_best.pth` - Best performing model (epoch 28) -- `taxonomy_embeddings_multi_groups.png` - Multi-group UMAP visualization -- `best_epoch_analysis_epoch28.png` - Comprehensive training analysis -- `umap_taxonomy_model_small_best_mammals_highlighted.png` - Single-group visualization - -**Performance:** -- **Loss:** 0.472317 (51.6% improvement) -- **Organisms:** 92,290 embedded in 10 dimensions -- **Ball Constraint:** 100% compliance (all embeddings inside unit ball) -- **Training Time:** ~2.5 hours on M3 Mac CPU - -**Quality Metrics:** -- Mean norm: 0.7147 -- Norm range: [0.0889, 1.0000] -- Clear hierarchical clustering -- Proper taxonomic group separation - -### **2. Reference Model: Animals Dataset (1M organisms)** - -**Location:** `taxonomy_model_animals_best.pth` - -**Performance:** -- **Loss:** 0.634712 (20.4% improvement after 4 epochs) -- **Organisms:** 1,055,469 embedded -- **Status:** Incomplete (needs 20+ more epochs for convergence) -- **Purpose:** Proof of scalability - ---- - -## 🔧 Core Components - -### **Training Pipeline** - -**Main Script:** `train_small.py` -- Depth-aware initialization -- Real-time metrics visualization -- Fixed early stopping logic -- Perfect ball constraint enforcement -- Hard negative sampling with depth stratification - -**Features:** -- Transitive closure training (975K pairs from 100K edges) -- Radial regularization (λ=0.1) -- Ranking loss with margin (0.2) -- Gradient clipping and selective projection -- Automatic checkpoint management - -### **Data Preparation** - -**Scripts:** -- `prepare_taxonomy_data.py` - Extract edges from NCBI taxonomy -- `remap_edges.py` - Create continuous index mapping -- `build_transitive_closure.py` - Build ancestor-descendant pairs -- `scripts/validate_data.py` - Data quality validation - -**Data Files:** -- `data/taxonomy_edges_small_transitive.pkl` - Training data (975K pairs) -- `data/taxonomy_edges_small.mapping.tsv` - TaxID→index mapping -- `data/names.dmp`, `data/nodes.dmp` - NCBI taxonomy structure - -### **Visualization** - -**Scripts:** -- `visualize_multi_groups.py` - Multi-group UMAP (Euclidean baseline) -- `visualize_animals_hyperbolic.py` - Hyperbolic-aware UMAP (correct geometry) - -**Key Groups Visualized:** -- Mammals (422 organisms) -- Birds (3,187 organisms) -- Insects (11,200 organisms) -- Bacteria (18,584 organisms) -- Fungi (1,002 organisms) -- Plants (14,744 organisms) - ---- - -## 🎯 Key Achievements - -### **1. Fixed Critical Bugs** - -✅ **Early Stopping Bug** -- **Problem:** Comparing loss against updated value (always self) -- **Fix:** Save previous best before updating -- **Impact:** Allowed training to reach epoch 28 (vs premature stop at epoch 5) - -✅ **Data Quality Issues** -- Fixed TaxID header contamination -- Corrected mapping inconsistencies -- Added comprehensive validation - -✅ **Hyperbolic Geometry** -- **Problem:** Using Euclidean distance for hyperbolic embeddings -- **Fix:** Implemented proper Poincaré distance metric -- **Impact:** Correct visualization of hierarchical structure - -### **2. Optimizations** - -- **1000x faster** regularizer (vectorized operations) -- **30x faster** projection (selective updates) -- **Stable training** on M3 Mac CPU (no GPU needed) - -### **3. Hierarchical Features** - -- Depth-aware initialization -- Hard negative sampling -- Depth-weighted loss -- Transitive closure (all ancestor-descendant pairs) - ---- - -## 📊 Results Summary - -### **Small Dataset (Recommended)** - -| Metric | Value | -|--------|-------| -| **Organisms** | 92,290 | -| **Training Pairs** | 975,131 | -| **Best Epoch** | 28 | -| **Loss** | 0.472 | -| **Improvement** | 51.6% | -| **Training Time** | 2.5 hours | -| **Model Size** | 3.5 MB | -| **Ball Constraint** | 100% ✅ | - -### **Animals Dataset (Incomplete)** - -| Metric | Value | -|--------|-------| -| **Organisms** | 1,055,469 | -| **Training Pairs** | 22,135,131 | -| **Epochs Trained** | 4 | -| **Loss** | 0.635 | -| **Improvement** | 20.4% | -| **Model Size** | 40.3 MB | -| **Status** | Needs more training | - ---- - -## 🚀 Usage - -### **Quick Start** - -```bash -# Train on small dataset (recommended) -python train_small.py - -# Visualize results -python visualize_multi_groups.py small_model_28epoch/taxonomy_model_small_best.pth - -# Check model quality -python check_model.py -``` - -### **Loading Embeddings** - -```python -import torch - -# Load model -ckpt = torch.load('small_model_28epoch/taxonomy_model_small_best.pth') -embeddings = ckpt['embeddings'] # Shape: (92290, 10) - -# Load TaxID mapping -import pandas as pd -mapping = pd.read_csv('data/taxonomy_edges_small.mapping.tsv', - sep='\t', header=None, names=['idx', 'taxid']) -``` - -### **Computing Distances** - -```python -import numpy as np - -def poincare_distance(x, y): - """Compute Poincaré distance (hyperbolic geometry).""" - diff_sq = np.sum((x - y)**2) - x_norm_sq = np.sum(x**2) - y_norm_sq = np.sum(y**2) - - ratio = 1 + 2 * diff_sq / ((1 - x_norm_sq) * (1 - y_norm_sq)) - return np.arccosh(np.clip(ratio, 1.0, None)) -``` - ---- - -## 🔬 Technical Insights - -### **What Worked** - -1. **Transitive closure** - Training on ALL ancestor-descendant pairs (not just parent-child) -2. **Depth-aware initialization** - Initialize embeddings by depth (shallow=center, deep=boundary) -3. **Radial regularization** - Strong regularization (λ=0.1) keeps embeddings in ball -4. **Long training** - 28 epochs needed for convergence (not 2-5) -5. **Small datasets** - 111K organisms is sweet spot for CPU training - -### **What Didn't Work** - -1. **Full dataset (2.7M)** - O(n²) complexity for hard negatives caused OOM -2. **Hard negatives at scale** - Sibling map construction doesn't scale beyond ~100K nodes -3. **Euclidean UMAP** - Wrong distance metric distorts hyperbolic structure -4. **Short training** - 4 epochs insufficient for hierarchy learning - -### **Scaling Challenges** - -- **Hard negatives:** O(n²) sibling map = ~7.3 trillion ops for 2.7M nodes -- **Poincaré distance:** O(n²) pairwise distances limits UMAP sample size -- **Solution:** Use random negatives + longer training for large datasets - ---- - -## 📝 Documentation - -### **Core Documents** - -- `README.md` - Project overview and installation -- `JOURNEY.md` - Complete development history -- `TRAIN_SMALL_GUIDE.md` - Training instructions -- `QUICKSTART.md` - Quick reference - -### **Code Structure** - -``` -poincare-embeddings/ -├── train_small.py # Main training script ⭐ -├── train_hierarchical.py # Core hierarchical model -├── visualize_multi_groups.py # UMAP visualization -├── build_transitive_closure.py # Data preparation -├── small_model_28epoch/ # Production model ⭐ -│ ├── taxonomy_model_small_best.pth -│ ├── taxonomy_embeddings_multi_groups.png -│ └── best_epoch_analysis_epoch28.png -├── data/ # NCBI taxonomy data -└── scripts/ # Utilities -``` - ---- - -## 🎓 Key Learnings - -### **1. Convergence Time is Critical** -- Small dataset needed 28 epochs (not 2-5) -- Early stopping must be implemented correctly -- Monitor metrics carefully - premature stopping ruins quality - -### **2. Hyperbolic Geometry Must Be Respected** -- Euclidean distance is wrong for hyperbolic embeddings -- Use Poincaré distance: `d = arcosh(1 + 2||x-y||²/((1-||x||²)(1-||y||²)))` -- Visualization requires proper metric - -### **3. Scaling Requires Different Strategies** -- Hard negatives don't scale (O(n²)) -- Use random negatives for large datasets -- Sample strategically for visualization (O(n²) distance computation) - -### **4. Training Stability is Paramount** -- Selective projection (only updated nodes) -- Periodic full projection (every 500 batches) -- Gradient clipping (max norm = 1.0) -- Strong regularization (λ = 0.1) - -### **5. Data Quality Matters** -- Fixed TaxID header bugs -- Comprehensive validation -- Proper mapping files - ---- - -## ✅ Production Checklist - -- [x] Training pipeline validated -- [x] Model converged (epoch 28) -- [x] Ball constraint enforced (100%) -- [x] Hierarchical structure verified -- [x] Visualizations generated -- [x] Documentation complete -- [x] Code cleaned and organized -- [x] Bugs fixed (early stopping, geometry) -- [x] Ready for downstream tasks - ---- - -## 🔮 Future Work (Optional) - -### **For Full Dataset** - -1. **Implement sparse hard negatives** - Don't store full sibling map -2. **Use approximate methods** - LSH or tree-based sampling -3. **GPU acceleration** - Batch Poincaré distance on GPU -4. **Incremental training** - Train in chunks, merge embeddings - -### **Model Improvements** - -1. **Riemannian optimizer** - Respects manifold natively -2. **Curriculum learning** - Start with parent-child, add deeper pairs -3. **Adaptive regularization** - Vary λ by depth level -4. **Alternative manifolds** - Try Lorentz or Klein models - -### **Applications** - -1. **Taxonomic prediction** - Predict parent given child -2. **Hierarchical clustering** - Group by hyperbolic distance -3. **Nearest neighbor queries** - Find related organisms -4. **Transfer learning** - Use embeddings as features - ---- - -## 📧 Contact - -For questions about this implementation, see the main README or open an issue on GitHub. - ---- - -## 🏆 Success Metrics - -| Goal | Status | Evidence | -|------|--------|----------| -| Train hierarchical embeddings | ✅ Complete | 92K organisms, loss 0.472 | -| Enforce ball constraint | ✅ Complete | 100% inside ball | -| Visualize hierarchy | ✅ Complete | Clear UMAP clusters | -| Fix critical bugs | ✅ Complete | Early stopping, geometry | -| Document process | ✅ Complete | JOURNEY.md, guides | -| Production ready | ✅ Complete | `small_model_28epoch/` | - ---- - -**Status: PRODUCTION READY** 🚀 - -The small dataset model is ready for deployment and downstream tasks. The codebase is clean, documented, and validated. diff --git a/docs/JOURNEY.md b/docs/JOURNEY.md deleted file mode 100644 index 753fec9..0000000 --- a/docs/JOURNEY.md +++ /dev/null @@ -1,904 +0,0 @@ -# Development Journey: From Facebook's Poincaré Embeddings to Hierarchical Taxonomy Embeddings - -## Overview - -This document chronicles the evolution of this project from Facebook Research's original Poincaré embeddings implementation to a specialized hierarchical taxonomy embedding system for NCBI's biological taxonomy. - ---- - -## Phase 1: Foundation - Facebook's Poincaré Embeddings - -### **Starting Point** -We began with Facebook Research's implementation of Poincaré embeddings, as described in their 2017 paper "Poincaré Embeddings for Learning Hierarchical Representations." - -**Original Features:** -- Hyperbolic geometry (Poincaré ball model) -- Designed for hierarchical data (e.g., WordNet) -- Pure parent-child edge training -- Simple negative sampling - -**Initial Success:** -- Successfully trained 2.7M NCBI taxonomy organisms -- Model converged in ~1 hour on CPU -- Basic hierarchical structure preserved - -### **Key Insight** -The Facebook implementation worked but was designed for simpler hierarchies. We needed enhancements for biological taxonomy with 38 depth levels and complex relationships. - ---- - -## Phase 2: Data Quality Issues (Nov 5-7, 2025) - -### **Critical Bugs Discovered** - -#### **Bug #1: Header Lines in Edgelist Files** -``` -Problem: .edgelist files had "id1 id2" header treated as real organism IDs -Result: 111,105 nodes (2 fake + 111,103 real) -Fix: Updated prepare_taxonomy_data.py to skip headers -Impact: Clean datasets with correct node counts -``` - -#### **Bug #2: Mapping File Inconsistencies** -``` -Problem: .mapping.tsv had fake TaxIDs "id1" and "id2" at indices 0 and 1 -Fix: Regenerated mappings without header artifacts -Impact: Proper TaxID → index mapping -``` - -#### **Bug #3: No Data Validation** -``` -Solution: Created scripts/validate_data.py -Features: - - Validates node counts - - Checks for duplicates - - Verifies index continuity - - Reports statistics -``` - -**Files Created:** -- `scripts/validate_data.py` -- `scripts/regenerate_data.sh` -- `DATA_FIXES_SUMMARY.md` - ---- - -## Phase 3: Hierarchical Training Attempt #1 (Nov 7-8, 2025) - -### **Motivation** -Simple parent-child training doesn't capture full hierarchy. Organisms at depth 20 should be "far" from root in hyperbolic space. - -### **Strategy: Transitive Closure** -Instead of 100K parent-child edges, train on ALL ancestor-descendant pairs: - -``` -Parent-child: 58,663 pairs (6%) -Grandparent: 52,953 pairs (5%) -Deep ancestors: 864,280 pairs (89%) -Total: 975,896 pairs (9.8x more data!) -``` - -### **Implementation** -Created `build_transitive_closure.py`: -- Loads NCBI taxonomy tree -- Computes all ancestor paths -- Adds depth metadata -- Outputs 975K training pairs - -### **Critical Bug: TaxID as Index** -``` -Problem: Used TaxIDs directly as embedding indices -Result: Created 3,467,244 embeddings (instead of 111,103) -Symptom: 97% of embeddings never updated! -Fix: Properly read mapping file with header -Impact: Reduced to 92,290 embeddings (correct) -``` - -**Files Created:** -- `build_transitive_closure.py` -- `BUGS_FOUND_AND_FIXED.md` -- `sanity_check.py` (comprehensive validation) - ---- - -## Phase 4: Hierarchical Features (Nov 8, 2025) - -### **New Training Features Implemented** - -#### **1. Depth-Aware Initialization** -```python -# Root (depth 0): radius ≈ 0.1 (near center) -# Leaves (depth 38): radius ≈ 0.95 (near boundary) -target_radius = 0.1 + (depth / max_depth) * 0.85 -``` - -#### **2. Radial Regularization** -```python -# Soft penalty to encourage depth → radius mapping -reg_loss = λ * (actual_radius - target_radius)² -``` - -#### **3. Hard Negative Sampling** -```python -# Sample negatives from same depth (cousins, not random) -siblings = nodes_at_same_depth(node) -negatives = random.sample(siblings, n_negatives) -``` - -#### **4. Depth Weighting** -```python -# Deeper pairs are more informative -weight = sqrt(depth_diff + 1) -loss = loss * weight -``` - -#### **5. Ranking Loss with Margin** -```python -# Encourage: d(ancestor, descendant) < d(ancestor, negative) + margin -loss = relu(pos_dist - neg_dist + margin) -``` - -**Files Created:** -- `train_hierarchical.py` -- `run_hierarchical_training.sh` - ---- - -## Phase 5: Ball Constraint Enforcement (Nov 8, 2025) - -### **The Problem: Embeddings Escaping** - -Training iterations showed embeddings violating the Poincaré ball constraint (||x|| ≥ 1.0): - -| Version | Max Norm | Outside Ball | Issue | -|---------|----------|--------------|-------| -| v1 | 2.18 | 50K (54%) | Weak regularization | -| v2 | 1.45 | 2K (2.2%) | Better but not enough | - -### **Root Causes** -1. Regularization too weak (λ=0.01) -2. Learning rate too high (0.01) -3. No gradient control -4. Projection insufficient - -### **Solution: 3-Layer Enforcement Strategy** - -#### **Layer 1: Improved Hyperparameters** -```python -Learning rate: 0.01 → 0.005 (2x slower) -Regularization: 0.01 → 0.1 (10x stronger) -Gradient clipping: None → max_norm=1.0 -``` - -#### **Layer 2: Hard Projection** -```python -# Only scale embeddings that violate constraint -needs_projection = (norms >= 1.0) -scale = where(needs_projection, (1-eps)/norms, 1.0) -``` - -#### **Layer 3: Periodic Full Projection** -```python -# Every 500 batches: project ALL embeddings -if n_batches % 500 == 0: - model.project_to_ball(indices=None) - -# End of epoch: GUARANTEE all inside ball -model.project_to_ball(indices=None) -``` - -### **Results** -| Version | Max Norm | Outside Ball | Status | -|---------|----------|--------------|--------| -| v3 | 1.00 | 0 (0%) | ✅ Perfect | - -**Files Created:** -- `BALL_CONSTRAINT_ENFORCEMENT.md` -- `TRAINING_OPTIMIZATIONS.md` - ---- - -## Phase 6: Performance Optimizations (Nov 8, 2025) - -### **Critical Optimizations** - -#### **1. Regularizer Vectorization** -``` -Before: Loop over 111K nodes per batch (1.7B ops/epoch) -After: Precompute tensors once (111K ops/epoch) -Speedup: 1000x -``` - -#### **2. Selective Projection** -``` -Before: Project all 92K embeddings every batch -After: Project only updated ~3K embeddings per batch -Speedup: 30x -``` - -#### **3. Efficient Tensor Creation** -``` -Before: List of numpy arrays → tensor (slow) -After: Pre-allocate numpy array → tensor (fast) -Speedup: 10-100x -``` - -#### **4. Device Selection** -``` -Before: Auto-select MPS on M3 Mac → hang -After: Force CPU (stable on macOS) -Result: Training actually runs! -``` - -**Overall:** From hanging indefinitely to ~3 minutes/epoch - ---- - -## Current State (Nov 8, 2025, 4:00pm) - -### **What Works ✅** - -1. **Data Pipeline** - - ✅ Correct TaxID → index mapping - - ✅ 92,290 embeddings (not 3.4M) - - ✅ 975,896 training pairs (transitive closure) - - ✅ Comprehensive validation (sanity_check.py) - -2. **Ball Constraints** - - ✅ 100% embeddings inside ball - - ✅ Max norm = 1.000 (perfect) - - ✅ 3-layer enforcement strategy - -3. **Training Stability** - - ✅ No crashes or hangs - - ✅ Gradient clipping prevents exploding gradients - - ✅ ~3 minutes per epoch on CPU - - ✅ Automatic checkpointing (keeps last 5 + best) - -### **What Doesn't Work ❌** - -**Hierarchy Quality is Poor:** -``` -Depth-norm correlation: +0.003 (target: >0.5) -Phylum separation: 1.08x (target: >1.5x) -Class separation: 0.99x (target: >1.5x) -Order separation: 0.99x (target: >1.5x) -``` - -**Hypothesis:** Only 2 epochs completed before manual stop. Possible issues: -- Need more training time -- Regularization too strong (λ=0.1 may be constraining) -- Data imbalance (94% deep pairs, only 6% parent-child) -- Margin too small (0.2) for hierarchical differences - ---- - -## Key Files & Scripts - -### **Core Training** -- `train_hierarchical.py` - Main hierarchical training script -- `build_transitive_closure.py` - Generate ancestor-descendant pairs -- `run_hierarchical_training.sh` - Quick-start training script - -### **Analysis** -- `analyze_hierarchy_hyperbolic.py` - Evaluate hierarchy quality -- `sanity_check.py` - Comprehensive validation (10 checks) - -### **Data Preparation** -- `prepare_taxonomy_data.py` - Download and prepare NCBI taxonomy -- `remap_edges.py` - Map TaxIDs to continuous indices -- `scripts/validate_data.py` - Validate data integrity - -### **Utilities** -- `scripts/visualize_embeddings.py` - UMAP visualization -- `watch_training.sh` - Monitor training progress -- `cleanup_for_release.sh` - Prepare for GitHub push - -### **Documentation** -- `README.md` - Main project documentation -- `QUICKSTART.md` - Get started in 5 minutes -- `JOURNEY.md` - This document -- `SESSION_SUMMARY_NOV8.md` - Latest session summary - ---- - -## Lessons Learned - -### **1. Data Quality is Critical** -- Always validate input data -- Don't assume file formats -- Use automated checks (sanity_check.py) -- A single wrong mapping can waste hours - -### **2. Constraints Are Hard** -- Poincaré ball constraint (||x|| < 1) requires careful handling -- Need multiple enforcement layers -- Trade-off between constraint and optimization freedom -- "Technically correct" ≠ "semantically good" - -### **3. Start Simple, Then Add Complexity** -- Facebook's simple model worked for 2.7M organisms -- Our complex model has perfect constraints but poor hierarchy -- Sometimes simpler is better -- Validate each feature independently - -### **4. Hyperbolic Space is Different** -- Distances grow exponentially near boundary -- Small norm differences = large distance differences -- Regularization interacts with geometry -- Need domain-specific tuning - -### **5. Documentation Matters** -- Intermediate docs helped track decisions -- Checkpoint management saves work -- Reproducibility requires clear instructions -- Future self will thank you - ---- - -## Next Steps & Open Questions - -### **Immediate Actions** -1. Train longer (try 20-50 epochs with patience=10) -2. If no improvement, reduce regularization (λ=0.05) -3. Monitor depth-norm correlation each epoch -4. Try balanced sampling (equal parent-child vs deep pairs) - -### **Research Questions** -1. Is transitive closure helping or hurting? -2. Are 31K siblings per node too many for hard negatives? -3. What's the optimal regularization strength? -4. Should margin vary by depth level? -5. Is curriculum learning needed (parent-child → grandparent → all)? - -### **Alternative Approaches** -1. Go back to simple training (known to work) -2. Use Riemannian optimizer (respects manifold natively) -3. Try different hyperbolic models (Lorentz, Klein) -4. Implement progressive training (build hierarchy bottom-up) - ---- - -## Technical Achievements - -### **Performance** -- ✅ 1000x faster regularizer (vectorized) -- ✅ 30x faster projection (selective) -- ✅ 100% ball constraint compliance -- ✅ Stable training on Mac M3 CPU - -### **Data Pipeline** -- ✅ Fixed critical TaxID bug -- ✅ Comprehensive validation suite -- ✅ Clean, reproducible data preparation -- ✅ Transitive closure computation - -### **Code Quality** -- ✅ Modular, well-documented code -- ✅ Automated testing (sanity_check.py) -- ✅ Professional repository structure -- ✅ Clear error messages and logging - ---- - -## Phase 6: Small Dataset Success (Nov 8-9, 2025) - -### **Breakthrough: Fixed Early Stopping Bug** - -**The Problem:** -```python -# BUG: tracker.update() updates best_loss BEFORE comparison -tracker.update(epoch, metrics) -if avg_loss < tracker.best_loss: # Always comparing against self! - epochs_without_improvement = 0 -``` - -**Result:** Early stopping triggered after 5 epochs despite continuous improvement. - -**The Fix:** -```python -# Save previous best BEFORE updating -prev_best_loss = tracker.best_loss -tracker.update(epoch, metrics) -if avg_loss < prev_best_loss: # Compare against OLD best - epochs_without_improvement = 0 -``` - -### **Training Success on Small Dataset (111K organisms)** - -Ran `train_small.py` with corrected early stopping: - -| Metric | Value | -|--------|-------| -| **Best Epoch** | 28 | -| **Final Loss** | 0.472317 | -| **Improvement** | 51.6% (from 0.977) | -| **Training Time** | ~2.5 hours | -| **Organisms** | 92,290 embedded | - -### **Model Quality Metrics** - -✅ **Perfect Ball Constraint** -- All embeddings inside unit ball (max norm = 1.0000) -- Mean norm: 0.7147 -- 24,979 nodes near boundary (deep in hierarchy) -- 9,275 nodes near center (root/shallow) - -✅ **Hierarchical Structure** -- Clear depth stratification -- Proper taxonomic clustering -- Meaningful nearest neighbors - -### **Visualization Results** - -**Multi-Group UMAP:** -- Mammals: 422 organisms -- Birds: 3,187 organisms -- Insects: 11,200 organisms -- Bacteria: 18,584 organisms -- Fungi: 1,002 organisms -- Plants: 14,744 organisms - -All groups properly clustered with clear separation! - ---- - -## Phase 7: Scaling Challenges (Nov 9-10, 2025) - -### **Attempt 1: Full Dataset (2.7M organisms)** - -**Problem:** Out of memory during transitive closure construction -```bash -python build_transitive_closure_full.py -# KILLED - Process terminated -``` - -**Root Cause:** O(n²) sibling map construction for hard negative sampling with 2.7M nodes = ~7.3 trillion operations - -### **Attempt 2: Animals Subset (Metazoa, 1.05M organisms)** - -**Strategy:** -1. Filter to animals only (TaxID 33208) -2. Build transitive closure: 22.1M training pairs -3. Use random negatives instead of hard negatives - -**Result:** Trained 4 epochs before manual stop - -| Metric | Value | -|--------|-------| -| **Organisms** | 1,055,469 | -| **Training Pairs** | 22,135,131 | -| **Best Loss** | 0.634712 (epoch 4) | -| **Improvement** | 20.4% | -| **Model Size** | 40.3 MB | - -### **Issue: Insufficient Training** -- Only 4 epochs vs 28 needed for convergence -- Loss still decreasing (not converged) -- UMAP showed scattered clusters - -**Created:** `continue_animals_training.py` to resume training (not used - opted to focus on small dataset success instead) - ---- - -## Phase 8: Hyperbolic Geometry Corrections (Nov 10, 2025) - -### **Critical Realization: Wrong Distance Metric** - -**The Problem:** -```python -# WRONG: Using Euclidean distance for hyperbolic embeddings -umap.UMAP(metric='euclidean') # ❌ Treats hyperbolic space as flat -``` - -Poincaré embeddings live in **hyperbolic space**, but we were visualizing them with **Euclidean distance** - fundamentally incorrect! - -### **The Fix: Proper Poincaré Distance** - -Implemented correct hyperbolic distance: -```python -def poincare_distance(x, y): - """ - Poincaré distance formula (respects hyperbolic geometry) - """ - diff_norm_sq = ||x - y||² - x_norm_sq = ||x||² - y_norm_sq = ||y||² - - ratio = 1 + 2 * diff_norm_sq / ((1 - x_norm_sq)(1 - y_norm_sq)) - return arcosh(ratio) - -# Correct UMAP usage -umap.UMAP(metric='precomputed') # Use precomputed Poincaré distances -``` - -### **Impact of Correction** - -| Visualization | Distance | Geometry | Result | -|--------------|----------|----------|--------| -| Previous | Euclidean | ❌ Wrong | Distorted, scattered | -| **Corrected** | **Poincaré** | **✅ Correct** | **True hierarchical structure** | - -**Key Difference:** -- Euclidean distance range: [0, ~2] -- Poincaré distance range: [0, ~19] - respects hyperbolic expansion - -### **Lessons Learned** - -1. **Geometry matters:** Hyperbolic embeddings require hyperbolic distances -2. **Complexity tradeoff:** Poincaré distance is O(n²), limiting sample size -3. **Validation importance:** Always verify mathematical correctness, not just implementation - ---- - -## Conclusion - -We've successfully built and validated a hierarchical Poincaré embedding system for biological taxonomy: - -### **✅ Achieved:** - -1. **Data Quality** - - Fixed critical TaxID bugs - - Comprehensive validation suite - - Transitive closure for hierarchy learning - -2. **Training Infrastructure** - - Fixed early stopping bug - - Real-time metrics visualization - - Depth-aware initialization - - Perfect ball constraint enforcement - -3. **Small Dataset Success (111K organisms)** - - **Best epoch: 28** - - **Loss: 0.472 (51% improvement)** - - **All embeddings inside ball** - - **Clear hierarchical clustering** - -4. **Mathematical Correctness** - - Proper Poincaré distance metric - - Hyperbolic-aware visualization - - Geometrically sound projections - -### **📊 Final Results:** - -**Small Model (Recommended):** -- 92,290 organisms embedded -- 3.5 MB model size -- Excellent hierarchical structure -- Ready for downstream tasks - -**Animals Model (Incomplete):** -- 1,055,469 organisms (4 epochs) -- Needs 20+ more epochs to converge -- Proof of scalability - -### **🎯 Key Insights:** - -1. **Convergence time is critical** - 28 epochs needed for quality hierarchy (not 2-5) -2. **Early stopping bugs are dangerous** - Can halt training prematurely -3. **Hyperbolic geometry must be respected** - Euclidean distance distorts structure -4. **Hard negatives don't scale** - O(n²) sibling maps fail beyond ~100K nodes -5. **Small datasets work beautifully** - 111K organisms is sweet spot for CPU training - -### **🚀 Production Ready:** - -The small model (`small_model_28epoch/`) is production-ready for: -- Taxonomic prediction -- Hierarchical queries -- Nearest neighbor search -- Downstream ML tasks - -### **📁 Repository State:** - -Clean, documented, and ready for deployment: -- Core training pipeline -- Validated data preparation -- Proper hyperbolic geometry -- Comprehensive documentation - ---- - -## Phase 9: Enhanced Training & Repository Organization (Nov 13-14, 2025) - -### **Training Script Evolution** - -#### **Three Training Approaches Now Available:** - -1. **`embed.py` - Original Facebook Trainer (Proven)** - - Battle-tested on 2.7M full dataset - - Uses simple `.mapped.edgelist` format - - Critical fix applied: embedding initialization scale = 0.1 (not 1e-4) - - Parameters validated from memory: `-lr 0.1 -burnin 10 -negs 50` - - Status: ✅ Fully working - -2. **`train_hierarchical.py` - Core Hierarchical Library** - - Implements advanced features: - - Depth-aware initialization - - Transitive closure training (975K pairs) - - Hard negative sampling (cousin nodes) - - Radial regularizer (depth → radius mapping) - - Proper Poincaré distance computation - - Early stopping with patience - - Status: ✅ Core implementation complete - -3. **`train_small.py` - Enhanced User Interface** - - Wrapper around `train_hierarchical.py` - - Enhanced terminal visualization: - - Color-coded improvements (green/red) - - Epoch-to-epoch comparisons (ΔLoss, % change) - - Real-time metrics (loss, reg, norm, outside%) - - Visual status indicators (✓ BETTER / ✗ WORSE) - - Better data handling (fills missing node depths) - - Automatic best model checkpointing - - Status: ✅ Production-ready for small dataset - -#### **Command for Extended Training:** -```bash -uv run python train_small.py --epochs 999999 --early-stopping 0 -``` -This runs indefinitely with enhanced visualization until manual stop (Ctrl+C). - -### **Repository Cleanup & Organization** - -#### **Archive Strategy Implemented:** - -**Moved to `docs/archive/debug_scripts/`:** -- `analyze_embeddings.py` - Early analysis attempts -- `analyze_messiness.py` - Hierarchy debugging -- `compare_old_new.py`, `compare_old_vs_current.py` - Comparison tools -- `diagnose_issues.py`, `find_what_broke.py` - Diagnostic scripts -- `inspect_checkpoint.py` - Simple checkpoint inspector -- `test_depth_coverage.py` - Test scripts -- `verify_ball_safety.py`, `verify_fixes.py` - Verification tools -- `visualize_trained_only.py` - Visualization variant - -**Moved to `docs/archive/`:** -- `PERMANENT_FIX_PLAN.md` - Old fix plans -- `PERMANENT_FIX_SUMMARY.md` - Fix summaries -- `REVERT_HYPERPARAMS.md` - Hyperparameter experiments -- `SAFETY_CHECK_BALL_CONSTRAINTS.md` - Constraint checks -- `TRAINING_ISSUES_FIXED.md` - Training notes - -#### **Final Repository Structure:** - -``` -Root Level (Core Files Only): -├── embed.py, train_hierarchical.py, train_small.py # 3 training approaches -├── prepare_taxonomy_data.py, build_transitive_closure.py # Data pipeline -├── analyze_hierarchy_hyperbolic.py, visualize_multi_groups.py # Analysis -├── sanity_check.py, final_sanity_check.py # Validation -├── README.md, QUICKSTART.md # Documentation - -Subdirectories: -├── hype/ # Original Facebook implementation -├── src/taxembed/ # New package structure (uv/pyproject.toml) -├── scripts/ # Utility scripts -├── tests/ # Unit tests -└── docs/ - ├── archive/ # Historical docs & debug scripts - ├── JOURNEY.md # This file - ├── FINAL_STATUS.md, CLI_COMMANDS.md, etc. -``` - -### **Current Status (Nov 14, 2025):** - -**✅ What Works:** -- Three complementary training approaches -- Clean, organized repository structure -- Comprehensive documentation -- All debugging tools archived (preserved for history) -- Enhanced visualization for training progress -- uv-based dependency management -- Production-ready small model (28 epochs) - -**🎯 Ready For:** -- Extended training runs (infinite epochs with early stop disabled) -- Public repository sharing -- Research collaboration -- Production deployment - -**📊 Key Metrics (Small Dataset - Best Model):** -- Organisms: 111,103 -- Training pairs: 975,000 (transitive closure) -- Epochs: 28 -- Loss: 0.472 -- Ball constraint compliance: 100% -- Visualization: Clear hierarchical clustering - -### **Lessons Learned:** - -1. **Multiple training approaches serve different purposes:** - - Original `embed.py` for validation and proven results - - `train_hierarchical.py` for core hierarchical features - - `train_small.py` for enhanced UX and monitoring - -2. **Preservation of development history is valuable:** - - Archived debug scripts document problem-solving journey - - Session notes capture decision rationale - - Future debugging benefits from traced history - -3. **Repository organization matters:** - - Clear separation: core scripts (root) vs utilities (scripts/) vs archives (docs/archive/) - - Gitignore properly excludes local artifacts (checkpoints, plots, data) - - Documentation structure supports different user needs - ---- - -## Phase 10: Unified CLI & Enhanced Visualization (Dec 2025) - -### **Streamlined Training & Visualization Pipeline** - -#### **The Vision** -Create a unified CLI that allows users to: -1. Train models for any clade with a single command: `taxembed train -as ` -2. Visualize results automatically: `taxembed visualize ` -3. Build custom datasets on-the-fly using TaxoPy -4. Track all artifacts (checkpoints, metadata, plots) in organized tag directories - -#### **Implementation: Unified CLI** - -**Created `src/taxembed/cli/main.py`:** -- Single entry point: `taxembed` command with `train` and `visualize` subcommands -- Automatic TaxID/clade name resolution using TaxoPy -- Dynamic dataset building via `taxopy_clade.py` builder -- Artifact management: all outputs stored in `artifacts/tags//` -- Metadata tracking: `run.json` stores training config, paths, dataset info - -**Key Features:** -```bash -# Train any clade by name or TaxID -taxembed train Cnidaria -as cnidaria --epochs 100 --lambda 0.1 -taxembed train 6073 -as echinoderms --epochs 50 - -# Visualize with automatic best checkpoint selection -taxembed visualize cnidaria -taxembed visualize echinoderms --children 1 # Color by grandchildren -``` - -#### **Dynamic Dataset Building** - -**Created `src/taxembed/builders/taxopy_clade.py`:** -- Queries NCBI taxonomy via TaxoPy for all descendants of a clade -- Builds parent-child edges automatically -- Computes transitive closure (all ancestor-descendant pairs) -- Remaps to sequential indices for efficient training -- Writes manifest with provenance (root TaxID, name, node counts, etc.) -- Progress bars for long-running operations - -**Benefits:** -- No manual data preparation needed -- Works for any taxonomic group -- Automatic handling of merged/obsolete TaxIDs -- Reproducible dataset generation - -#### **Enhanced Visualization** - -**Updated `visualize_multi_groups.py`:** -- Automatic best checkpoint selection per tag -- Hierarchical coloring: `--children` flag controls depth (0=children, 1=grandchildren, etc.) -- Informative titles: `"TaxEmbed: {CLADE}, Children Level {X}, epochs {Y}, Loss {L}"` -- Robust path resolution: works regardless of current working directory -- TaxoPy fallback: uses TaxoPy if local dump files are missing - -**Visualization Features:** -- UMAP dimensionality reduction with proper Poincaré distance -- Color-coded by taxonomic groups (children/grandchildren of root) -- Automatic legend with group names and counts -- High-quality output suitable for publications - -#### **Artifact Management** - -**Organized Structure:** -``` -artifacts/tags/ -├── cnidaria/ -│ ├── run.json # Metadata (config, paths, dataset info) -│ ├── cnidaria.pth # Checkpoints -│ ├── cnidaria_best.pth # Best checkpoint -│ └── cnidaria_umap.png # Visualizations -├── echinoderms/ -│ └── ... -└── mammals/ - └── ... -``` - -**Metadata (`run.json`) includes:** -- Tag and slug -- Creation timestamp -- Dataset info (root TaxID, name, node counts, paths) -- Training config (epochs, learning rate, regularization, etc.) -- Paths to all artifacts (checkpoints, mapping files, data) - -#### **Progress Feedback** - -**User Experience Improvements:** -- Progress bars during dataset building (transitive closure computation) -- Color-coded training output (green for improvements, red for regressions) -- Clear status messages at each stage -- Automatic checkpoint path resolution - -#### **Technical Achievements** - -1. **Robust Path Handling:** - - Scripts resolve paths relative to their own location - - Works regardless of current working directory - - Handles both absolute and relative checkpoint paths - -2. **Error Handling:** - - Graceful fallbacks (TaxoPy if dump files missing) - - Clear error messages with actionable guidance - - Validation of inputs before processing - -3. **Code Organization:** - - Shared utilities in `src/taxembed/utils/` - - Modular builders in `src/taxembed/builders/` - - Clean separation of concerns - -#### **Example Workflow** - -```bash -# 1. Train a model for Cnidaria (jellyfish, corals, etc.) -taxembed train Cnidaria -as cnidaria --epochs 100 --lambda 0.1 - -# This automatically: -# - Resolves "Cnidaria" to TaxID 6072 -# - Builds dataset with all descendants -# - Trains model with specified hyperparameters -# - Saves checkpoints and metadata to artifacts/tags/cnidaria/ - -# 2. Visualize results -taxembed visualize cnidaria --children 0 # Color by immediate children - -# 3. Try different coloring depths -taxembed visualize cnidaria --children 1 # Color by grandchildren -``` - -#### **Files Created/Modified** - -**New Files:** -- `src/taxembed/cli/main.py` - Unified CLI entry point -- `src/taxembed/builders/taxopy_clade.py` - Dynamic dataset builder -- `src/taxembed/utils/data_validation.py` - Shared validation utilities -- `scripts/build_clade_dataset.py` - Standalone dataset builder script - -**Modified Files:** -- `visualize_multi_groups.py` - Enhanced with hierarchical coloring and informative titles -- `train_small.py` - Updated to handle dynamic mapping files -- `pyproject.toml` - Added `taxopy` dependency and unified CLI entry point -- `README.md` - Updated with new CLI usage examples - -#### **Lessons Learned** - -1. **User Experience Matters:** - - Single command workflows are much better than multi-step processes - - Progress bars prevent perceived "hangs" during long operations - - Informative titles help users understand what they're looking at - -2. **Robustness is Critical:** - - Path resolution must work from any directory - - Fallbacks (TaxoPy) prevent failures when files are missing - - Clear error messages save debugging time - -3. **Metadata is Essential:** - - Storing run metadata enables automatic visualization - - Provenance tracking (dataset info) ensures reproducibility - - Organized artifact structure makes it easy to find results - -4. **Modular Design:** - - Shared utilities prevent code duplication - - Builders can be used standalone or via CLI - - Clear separation allows independent testing - ---- - -## References - -- Nickel & Kiela (2017). "Poincaré Embeddings for Learning Hierarchical Representations" -- Facebook Research: https://github.com/facebookresearch/poincare-embeddings -- NCBI Taxonomy: https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/ -- TaxoPy: https://pypi.org/project/taxopy/ -- This project: https://github.com/jcoludar/taxembed - ---- - -*Last Updated: December 2025* diff --git a/docs/PRE_PUSH_CHECKLIST.md b/docs/PRE_PUSH_CHECKLIST.md deleted file mode 100644 index afa0217..0000000 --- a/docs/PRE_PUSH_CHECKLIST.md +++ /dev/null @@ -1,180 +0,0 @@ -# Pre-Push Checklist - -## ✅ Repository Cleanup Complete! - -### **Files Cleaned** -- [x] Removed old checkpoints (kept v3_best only) -- [x] Removed training logs -- [x] Removed temporary plots (.png files) -- [x] Moved intermediate docs to docs/archive/ -- [x] Cleaned Python cache (__pycache__, *.pyc) - -### **Documentation Updated** -- [x] README.md - Reflects current state and features -- [x] JOURNEY.md - Complete development history created -- [x] QUICKSTART.md - Exists and is current -- [x] SESSION_SUMMARY_NOV8.md - Latest findings documented - -### **Core Files Present** -- [x] train_hierarchical.py - Main training script -- [x] build_transitive_closure.py - Data preparation -- [x] analyze_hierarchy_hyperbolic.py - Analysis tool -- [x] sanity_check.py - Validation suite -- [x] run_hierarchical_training.sh - Quick-start script - -### **Data Files** -- [x] data/ directory exists -- [x] .gitignore properly excludes data files -- [x] One checkpoint kept: taxonomy_model_hierarchical_small_v3_best.pth (but gitignored) - ---- - -## 📝 Git Commands - -### **1. Check Status** -```bash -git status -``` - -Expected: Many deleted files (checkpoints, logs, plots, docs) and updated files (README, JOURNEY) - -### **2. Add All Changes** -```bash -git add -A -``` - -### **3. Commit** -```bash -git commit -m "Major cleanup and documentation overhaul - -- Consolidated development history into JOURNEY.md -- Updated README to reflect current state -- Removed temporary files (checkpoints, logs, plots) -- Moved intermediate docs to docs/archive/ -- Kept only essential files and v3_best checkpoint - -Key features documented: -- Transitive closure training (975K pairs) -- Hierarchical features (depth-aware, hard negatives) -- 3-layer ball constraint enforcement (100% compliance) -- Performance optimizations (1000x faster) - -Current status: -- Technical implementation complete -- Hierarchy quality needs improvement (more training) -- See JOURNEY.md for full development history" -``` - -### **4. Push** -```bash -git push origin main -``` - ---- - -## 🔍 Final Checks - -### **Run Sanity Check** -```bash -python sanity_check.py -``` -Expected: 10/10 checks passed - -### **Verify Documentation Links** -- [ ] README.md links work (JOURNEY.md, QUICKSTART.md, etc.) -- [ ] Code examples in README are correct -- [ ] Installation instructions are accurate - -### **Test Quick Start** -```bash -# In a fresh virtual environment -python3.11 -m venv test_venv -source test_venv/bin/activate -pip install -r requirements.txt -python sanity_check.py # Should pass -``` - ---- - -## 📊 Repository Statistics - -### **Before Cleanup** -- ~20 checkpoint files (~70MB total) -- Multiple training logs -- ~10 temporary plot files -- ~20 intermediate documentation files - -### **After Cleanup** -- 1 checkpoint (taxonomy_model_hierarchical_small_v3_best.pth, gitignored) -- 0 logs -- 0 temporary plots -- 3 core docs + archived intermediate docs - -### **Size Reduction** -- Removed ~70MB of checkpoints -- Removed ~5MB of plots and logs -- Organized docs (moved, not deleted) - ---- - -## 🎯 What Users Will See - -### **Main Files** -``` -README.md - Complete project overview -QUICKSTART.md - 5-minute setup guide -JOURNEY.md - Full development history -CONTRIBUTING.md - How to contribute -CODE_OF_CONDUCT.md - Community guidelines -``` - -### **Core Scripts** -``` -train_hierarchical.py - Main training -build_transitive_closure.py - Data prep -analyze_hierarchy_hyperbolic.py - Analysis -sanity_check.py - Validation -``` - -### **What's Hidden** -``` -data/ - Gitignored (users download) -*.pth - Gitignored (users train) -*.log, *.png - Gitignored (temporary) -venv*/ - Gitignored (local env) -``` - ---- - -## ✨ Key Selling Points - -1. **Complete Development History** - JOURNEY.md documents every decision -2. **Professional Structure** - Clean, organized, documented -3. **Reproducible** - Sanity checks, clear instructions -4. **Honest Documentation** - Documents what works AND what doesn't -5. **Ready for Research** - Clear future directions outlined - ---- - -## 🚀 Post-Push TODO - -After pushing: -1. Add repository link to README.md -2. Create GitHub release for v3_best checkpoint -3. Add issues for known limitations -4. Create project board for future work -5. Consider adding CI/CD for sanity checks - ---- - -## ✅ Ready to Push! - -All cleanup complete. Repository is professional, documented, and ready for public viewing. - -**Last step:** Review git diff, then push! - -```bash -git diff --stat # Review what changed -git log -1 # Review commit message -git push origin main -``` diff --git a/docs/PROJECT_TREE.txt b/docs/PROJECT_TREE.txt deleted file mode 100644 index e7b8f54..0000000 --- a/docs/PROJECT_TREE.txt +++ /dev/null @@ -1,123 +0,0 @@ -taxembed/ -│ -├── 📁 src/ -│ └── 📁 taxembed/ # Main package (src/ layout) -│ ├── __init__.py # Package initialization -│ ├── 📁 manifolds/ # Hyperbolic manifold implementations -│ │ └── __init__.py -│ ├── 📁 models/ # Embedding models -│ │ └── __init__.py -│ ├── 📁 datasets/ # Data loading and processing -│ │ └── __init__.py -│ └── 📁 utils/ # Utility functions -│ └── __init__.py -│ -├── 📁 scripts/ # Standalone executable scripts -│ ├── train.py # Main training script -│ ├── prepare_data.py # Data preparation wrapper -│ ├── remap_data.py # ID remapping wrapper -│ ├── monitor.py # Training monitoring wrapper -│ ├── evaluate.py # Model evaluation wrapper -│ └── visualize.py # Visualization wrapper -│ -├── 📁 tests/ # Unit and integration tests -│ ├── __init__.py -│ └── test_example.py # Example test module -│ -├── 📁 data/ # Data directory (gitignored) -│ ├── taxonomy_edges.edgelist # Raw edge list -│ ├── taxonomy_edges.mapped.edgelist # Remapped edge list -│ └── taxonomy_edges.mapping.tsv # ID mapping -│ -├── 📁 hype/ # Original package (backward compatibility) -│ ├── __init__.py -│ ├── manifolds/ -│ ├── models/ -│ └── ... (original structure) -│ -├── 📄 pyproject.toml # Project configuration (uv) -├── 📄 ruff.toml # Ruff linter configuration -├── 📄 Makefile # Convenient command shortcuts -├── 📄 setup.py # Setup script for C++ extensions -│ -├── 📄 README.md # Main project documentation -├── 📄 QUICKSTART.md # Quick start guide -├── 📄 GETTING_STARTED.md # Getting started guide -├── 📄 STRUCTURE.md # Project organization documentation -├── 📄 CONTRIBUTING.md # Contribution guidelines -├── 📄 RESTRUCTURING_SUMMARY.md # Migration guide -├── 📄 RESTRUCTURING_COMPLETE.md # Restructuring summary -├── 📄 PROJECT_TREE.txt # This file -│ -├── 📄 LICENSE # CC-BY-NC 4.0 license -├── 📄 .gitignore # Git ignore rules -│ -└── 📁 .git/ # Git repository - -═══════════════════════════════════════════════════════════════════════════════ - -📊 STATISTICS - -Files Created: - - Configuration: 3 (pyproject.toml, ruff.toml, Makefile) - - Documentation: 7 (README.md, QUICKSTART.md, GETTING_STARTED.md, etc.) - - Source Code: 5 (__init__.py files in src/taxembed/) - - Scripts: 6 (train.py, prepare_data.py, remap_data.py, etc.) - - Tests: 2 (__init__.py, test_example.py) - - Other: 2 (.gitignore, PROJECT_TREE.txt) - -Total New Files: 25+ - -═══════════════════════════════════════════════════════════════════════════════ - -🎯 KEY DIRECTORIES - -src/taxembed/ → Main package code (src/ layout best practice) -scripts/ → Standalone executable scripts -tests/ → Unit and integration tests -data/ → Data files (gitignored) -hype/ → Original package (backward compatible) - -═══════════════════════════════════════════════════════════════════════════════ - -📚 DOCUMENTATION HIERARCHY - -START HERE: - └─ GETTING_STARTED.md (5-minute setup) - ├─ QUICKSTART.md (detailed quick start) - ├─ README.md (full documentation) - ├─ STRUCTURE.md (project organization) - ├─ CONTRIBUTING.md (development guidelines) - └─ RESTRUCTURING_SUMMARY.md (what changed) - -═══════════════════════════════════════════════════════════════════════════════ - -⚙️ COMMON COMMANDS - -make help → Show all available commands -make install → Install dependencies with uv -make build → Build C++ extensions -make lint → Check code with ruff -make format → Format code with ruff -make test → Run tests with pytest -make test-cov → Run tests with coverage -make clean → Clean build artifacts - -uv run python scripts/train.py --help -uv run python scripts/prepare_data.py -uv run python scripts/evaluate.py --checkpoint model.pth - -═══════════════════════════════════════════════════════════════════════════════ - -✨ FEATURES - -✅ Professional Python project layout (src/ layout) -✅ Fast dependency management with uv -✅ Code quality enforcement with ruff -✅ Comprehensive documentation -✅ Testing framework (pytest) -✅ Convenient Makefile commands -✅ Backward compatibility maintained -✅ IDE-ready with type hints support - -═══════════════════════════════════════════════════════════════════════════════ diff --git a/docs/RELEASE_SUMMARY.md b/docs/RELEASE_SUMMARY.md deleted file mode 100644 index 575834a..0000000 --- a/docs/RELEASE_SUMMARY.md +++ /dev/null @@ -1,274 +0,0 @@ -# Release Summary - November 8, 2025 - -## 🎉 Repository Ready for GitHub Push! - -### **What We Did Today** - -1. ✅ **Fixed Critical Bugs** - - TaxID mapping bug (3.4M → 92K embeddings) - - Ball constraint violations (100% compliance achieved) - - Data validation pipeline - -2. ✅ **Implemented Hierarchical Features** - - Transitive closure training (975K pairs) - - Depth-aware initialization and regularization - - Hard negative sampling (cousins at same depth) - - 3-layer ball constraint enforcement - -3. ✅ **Performance Optimizations** - - 1000x faster regularizer (vectorized) - - 30x faster projection (selective) - - Stable training on M3 Mac CPU (~3 min/epoch) - -4. ✅ **Complete Documentation** - - README.md - Project overview - - JOURNEY.md - Development history - - QUICKSTART.md - Setup guide - - PRE_PUSH_CHECKLIST.md - Verification steps - -5. ✅ **Repository Cleanup** - - Removed 70MB+ of checkpoints - - Archived intermediate documentation - - Cleaned temporary files - - Professional structure - ---- - -## 📦 What's Being Committed - -### **New Files** (29 files) -``` -✅ Core Documentation - - JOURNEY.md (development history) - - PRE_PUSH_CHECKLIST.md (verification) - - RELEASE_SUMMARY.md (this file) - -✅ Core Scripts - - train_hierarchical.py (main training) - - build_transitive_closure.py (data prep) - - analyze_hierarchy_hyperbolic.py (analysis) - - sanity_check.py (validation) - - run_hierarchical_training.sh (quick-start) - -✅ Utility Scripts - - cleanup_for_release.sh - - watch_training.sh - - Various analysis scripts - -✅ Archive - - docs/archive/ (intermediate docs) -``` - -### **Modified Files** (2 files) -``` -📝 README.md - Complete rewrite reflecting current state -📝 scripts/visualize_embeddings.py - Updated -``` - -### **Deleted Files** (13 files) -``` -🗑️ Intermediate documentation (moved to docs/archive/) - - CLEANUP_SUMMARY.md - - DATA_FIXES_SUMMARY.md - - REPOSITORY_STATUS.md - - STRUCTURE.md - - ... and 9 more -``` - -### **Ignored Files** (not tracked) -``` -🚫 Checkpoints (.pth files) -🚫 Training logs (.log files) -🚫 Plots (.png files) -🚫 Data directory -🚫 Virtual environments -``` - ---- - -## 📊 Repository Statistics - -### **Code** -- **Python scripts:** 15+ core files -- **Lines of code:** ~5,000+ (train_hierarchical.py, sanity_check.py, etc.) -- **Documentation:** 3 core docs + archived intermediates -- **Tests:** Comprehensive sanity_check.py (10 checks) - -### **Features** -- ✅ Transitive closure computation -- ✅ Hierarchical training (5 advanced features) -- ✅ Ball constraint enforcement (3 layers) -- ✅ Performance optimizations (1000x speedups) -- ✅ Comprehensive validation -- ✅ Analysis and visualization tools - -### **Quality** -- ✅ All code documented -- ✅ Validated with sanity checks -- ✅ Professional structure -- ✅ Honest about limitations -- ✅ Clear future directions - ---- - -## 🎯 Current State - -### **Technical Achievements ✅** -- **Data Pipeline:** Clean, validated, reproducible -- **Training:** Stable, fast (~3 min/epoch), checkpointed -- **Constraints:** 100% embeddings inside Poincaré ball -- **Code Quality:** Modular, documented, tested - -### **Research Status ⚠️** -- **Hierarchy Quality:** Poor after limited training (2 epochs) -- **Depth Correlation:** ~0 (needs improvement) -- **Taxonomic Separation:** <1.1x (needs improvement) - -**Honest Assessment:** Technical implementation is solid, but hierarchy learning needs more work (tuning or training time). - ---- - -## 🚀 Ready to Push! - -### **Pre-Push Command** -```bash -# Review changes -git status -git diff --stat - -# Run final validation -python sanity_check.py # Should pass 10/10 - -# Add all changes -git add -A - -# Commit with detailed message -git commit -m "Major cleanup and documentation overhaul - -- Consolidated development history into JOURNEY.md -- Updated README to reflect current state -- Removed temporary files (checkpoints, logs, plots) -- Moved intermediate docs to docs/archive/ -- Kept only essential files - -Key features documented: -- Transitive closure training (975K pairs) -- Hierarchical features (depth-aware, hard negatives) -- 3-layer ball constraint enforcement -- Performance optimizations (1000x faster) - -Current status: -- Technical implementation complete -- Hierarchy quality needs improvement -- See JOURNEY.md for full development history" - -# Push to GitHub -git push origin main -``` - ---- - -## 📝 Post-Push TODO - -After pushing to GitHub: - -1. **Update README links** ✅ Done! - - ✅ Replaced `[Your Name]` with @jcoludar - - ✅ Updated repository to jcoludar/taxembed - -2. **Create Release** - - Tag: v0.1.0-alpha - - Title: "Initial Public Release - Technical Implementation" - - Attach: taxonomy_model_hierarchical_small_v3_best.pth - -3. **GitHub Issues** - - "Improve hierarchy quality (depth correlation ~0)" - - "Implement balanced sampling strategy" - - "Test on full 2.7M organism dataset" - - "Add curriculum learning" - -4. **Project Board** - - TODO: Hyperparameter tuning - - TODO: Longer training experiments - - TODO: Alternative sampling strategies - - DONE: Data pipeline - - DONE: Ball constraints - -5. **CI/CD** (optional) - - GitHub Actions for sanity_check.py - - Automated testing on push - - Documentation deployment - ---- - -## 🌟 What Makes This Repository Special - -### **1. Complete Transparency** -- Documents what works AND what doesn't -- Full development history (JOURNEY.md) -- Honest about current limitations - -### **2. Research-Ready** -- Comprehensive validation tools -- Clear experimental results -- Well-defined future directions - -### **3. Professional Quality** -- Clean, organized structure -- Thorough documentation -- Reproducible results - -### **4. Learning Resource** -- Shows real development process -- Documents debugging journey -- Explains design decisions - ---- - -## 📚 Key Documents - -For new users, read in this order: - -1. **README.md** - What is this project? -2. **QUICKSTART.md** - How do I use it? -3. **JOURNEY.md** - How did we get here? -4. **PRE_PUSH_CHECKLIST.md** - What was cleaned up? - -For contributors: - -1. **CONTRIBUTING.md** - How to contribute -2. **CODE_OF_CONDUCT.md** - Community guidelines -3. **docs/archive/** - Detailed development notes - ---- - -## ✨ Final Thoughts - -This repository represents a significant effort to extend Facebook's Poincaré embeddings for biological taxonomy. While the technical implementation is solid (perfect ball constraints, optimized performance, comprehensive validation), the hierarchy learning still needs work. - -**This is a great starting point for research**, with all the hard engineering problems solved: -- ✅ Data pipeline -- ✅ Efficient training -- ✅ Constraint enforcement -- ✅ Validation tools - -**The remaining challenges are research questions:** -- How to best sample training data? -- What's the optimal regularization strength? -- Should we use curriculum learning? - -We've built a solid foundation. Now it's time to experiment! - ---- - -## 🙏 Acknowledgments - -- Facebook Research for the original Poincaré embeddings implementation -- NCBI for the taxonomic data -- All the intermediate documentation that helped track our journey - ---- - -**Ready to share with the world! 🚀** - -*Generated: November 8, 2025* diff --git a/docs/TRAIN_FULL_GUIDE.md b/docs/TRAIN_FULL_GUIDE.md deleted file mode 100644 index be84162..0000000 --- a/docs/TRAIN_FULL_GUIDE.md +++ /dev/null @@ -1,229 +0,0 @@ -# Training on Full Dataset - Complete Guide - -## Overview - -Train hierarchical Poincaré embeddings on the **full NCBI taxonomy** (2.7M organisms) using the same successful approach from the small dataset. - ---- - -## Step 1: Build Transitive Closure - -**⏱️ Time Required: 30-60 minutes** - -```bash -python build_transitive_closure_full.py -``` - -### What This Does: -- Loads full NCBI taxonomy (2.7M+ organisms) -- Computes all ancestor-descendant pairs (not just parent-child) -- Creates training data with depth metadata -- Generates ~100M+ training pairs - -### Output Files: -- `data/taxonomy_edges_transitive.pkl` - Training data with metadata -- `data/taxonomy_edges_transitive.tsv` - Human-readable format -- `data/taxonomy_edges_transitive.edgelist` - Edge list format - ---- - -## Step 2: Train Model - -**⏱️ Time Required: 3-5 hours on M3 Mac CPU** - -```bash -python train_full.py -``` - -### Default Configuration: -Based on successful small dataset training: - -| Parameter | Value | Notes | -|-----------|-------|-------| -| Dimensions | 10 | Same as small dataset | -| Batch Size | 128 | Larger for efficiency | -| Learning Rate | 0.005 | Proven effective | -| Margin | 0.2 | Ranking loss margin | -| Regularization (λ) | 0.1 | Ball constraint | -| Negative Samples | 50 | Hard negatives | -| Early Stopping | 10 epochs | Higher patience for large dataset | -| Max Epochs | 100 | Will likely stop earlier | - -### What You'll See: - -``` -====================================================================================================== - Epoch | Loss | ΔLoss | Improve | Reg | MaxNorm | Outside | Status -====================================================================================================== -Epoch 1/100: 45%|████████████▌ | 6845/15248 batches [05:23<06:42] - 1 | 0.987654 | --- | --- | 0.012345 | 1.0000 | 0.00% | FIRST - └─ Best: 0.987654 @ epoch 1 | Norms: [0.0889, 0.6147, 1.0000] -``` - ---- - -## Step 3: Monitor Training - -Training will automatically: -- ✅ Display real-time progress bar per epoch -- ✅ Show loss improvements with color coding -- ✅ Save last 5 epoch checkpoints -- ✅ Save best model automatically -- ✅ Stop early when no improvement - -### Expected Timeline: -- Epoch 1-5: ~5-10 min/epoch (initial learning) -- Epoch 6-20: ~5-10 min/epoch (convergence) -- Epoch 20+: Likely early stopping - ---- - -## Parameter Tuning - -### For Faster Training (Less Accuracy): -```bash -python train_full.py \ - --batch-size 256 \ - --n-negatives 30 \ - --early-stopping 5 -``` - -### For Better Quality (Slower): -```bash -python train_full.py \ - --batch-size 64 \ - --n-negatives 100 \ - --early-stopping 15 \ - --lr 0.003 -``` - -### For Different Checkpoint Name: -```bash -python train_full.py --checkpoint taxonomy_model_full_v2.pth -``` - ---- - -## Output Files - -After training: -- `taxonomy_model_full_best.pth` - Best model (lowest loss) -- `taxonomy_model_full_epoch{N}.pth` - Last 5 epoch checkpoints -- `taxonomy_model_full.pth` - Final model - -Each checkpoint contains: -- Embeddings (2.7M × 10 dimensions) -- Training metadata (epoch, loss, etc.) -- Model state for resuming - ---- - -## Comparison: Small vs Full - -| Metric | Small Dataset | Full Dataset | -|--------|--------------|--------------| -| Organisms | 111K | 2.7M | -| Training Pairs | 975K | ~100M+ | -| Model Size | 3.5 MB | ~210 MB | -| Training Time | 2-3 hours | 3-5 hours | -| Epoch Time | ~3 min | ~5-10 min | -| Best Epoch | 28 | TBD | -| Final Loss | 0.472 | TBD | - ---- - -## Troubleshooting - -### "Training data not found" -Run `python build_transitive_closure_full.py` first. This takes 30-60 minutes. - -### Out of Memory -```bash -python train_full.py --batch-size 64 -``` - -### Training Too Slow -```bash -python train_full.py --batch-size 256 --n-negatives 30 -``` - -### Want to Use GPU -```bash -python train_full.py --gpu 0 -``` -(Requires CUDA-enabled GPU) - ---- - -## After Training - -### 1. Check Results -```bash -python check_model.py -``` - -### 2. Plot Analysis -```bash -python plot_best_epoch.py -``` - -### 3. Visualize Embeddings -```bash -python visualize_multi_groups.py taxonomy_model_full_best.pth -``` - -### 4. Analyze Hierarchy -```bash -python analyze_hierarchy_hyperbolic.py -``` - ---- - -## Expected Results - -Based on small dataset success, the full model should achieve: -- ✅ Loss reduction: ~50% improvement (0.98 → 0.47) -- ✅ All embeddings inside ball (0% outside) -- ✅ Proper hierarchical clustering -- ✅ Meaningful nearest neighbors -- ✅ Clear taxonomic group separation - ---- - -## Key Differences from Small Dataset - -1. **Batch Size**: 128 (vs 64) - larger for efficiency -2. **Early Stopping**: 10 epochs (vs 5) - more patience for complex dataset -3. **Training Time**: 3-5 hours (vs 2-3 hours) -4. **Model Size**: 210 MB (vs 3.5 MB) - -All other parameters are identical to the successful small dataset training! - ---- - -## Pro Tips - -1. **Run overnight**: Full training takes 3-5 hours -2. **Monitor early epochs**: If loss doesn't decrease in first 3 epochs, something's wrong -3. **Save intermediate checkpoints**: Last 5 epochs auto-saved -4. **Check max norm**: Should always be ≤ 1.0 (ball constraint) -5. **Compare to small**: Similar convergence pattern expected - ---- - -## Resume Training (If Interrupted) - -Training will need to be restarted from scratch. The current implementation doesn't support resuming. Consider: -- Using a screen/tmux session -- Running overnight when uninterrupted -- Monitoring the first few epochs to ensure proper training - ---- - -## Next Steps After Full Training - -1. Compare full vs small model performance -2. Test on downstream tasks (e.g., taxonomic prediction) -3. Evaluate embedding quality metrics -4. Visualize major taxonomic groups -5. Query nearest neighbors for validation diff --git a/docs/TRAIN_SMALL_GUIDE.md b/docs/TRAIN_SMALL_GUIDE.md deleted file mode 100644 index c6dbb37..0000000 --- a/docs/TRAIN_SMALL_GUIDE.md +++ /dev/null @@ -1,128 +0,0 @@ -# Training on Small Dataset - Quick Guide - -## Quick Start - -```bash -python train_small.py -``` - -That's it! The script is pre-configured for the small dataset with optimal defaults. - -## What You'll See - -The script displays a **real-time metrics table** showing: - -``` -====================================================================================================== - Epoch | Loss | ΔLoss | Improve | Reg | MaxNorm | Outside | Status -====================================================================================================== - 1 | 0.234567 | --- | --- | 0.045678 | 0.9876 | 0.00% | FIRST - 2 | 0.198765 | -0.035802 | -15.27% | 0.043210 | 0.9654 | 0.00% | ✓ BETTER - 3 | 0.187654 | -0.011111 | -5.59% | 0.041234 | 0.9543 | 0.00% | ✓ BETTER - 4 | 0.195432 | +0.007778 | +4.14% | 0.042567 | 0.9678 | 0.00% | ✗ WORSE -``` - -### Columns Explained: -- **Epoch**: Current epoch number -- **Loss**: Training loss for this epoch -- **ΔLoss**: Change from previous epoch (negative = better) -- **Improve**: Percentage improvement (green ✓ = better, red ✗ = worse) -- **Reg**: Regularization loss (keeps embeddings in ball) -- **MaxNorm**: Maximum embedding norm (should be < 1.0) -- **Outside**: Percentage of embeddings outside ball (should be 0%) -- **Status**: Visual indicator of improvement - -## Custom Parameters - -### Faster Training (Fewer Epochs) -```bash -python train_small.py --epochs 50 --early-stopping 3 -``` - -### Stronger Regularization (Keep embeddings tighter in ball) -```bash -python train_small.py --lambda-reg 0.2 -``` - -### Larger Batches (Faster, less precise) -```bash -python train_small.py --batch-size 128 -``` - -### Higher Learning Rate (Faster convergence, less stable) -```bash -python train_small.py --lr 0.01 -``` - -### Save to Custom Location -```bash -python train_small.py --checkpoint models/my_model.pth -``` - -## Full Parameter List - -```bash -python train_small.py --help -``` - -Available options: -- `--data`: Training data path (default: small transitive closure) -- `--checkpoint`: Output model path -- `--dim`: Embedding dimension (default: 10) -- `--epochs`: Maximum epochs (default: 100) -- `--early-stopping`: Patience before stopping (default: 5) -- `--batch-size`: Batch size (default: 64) -- `--n-negatives`: Negative samples per positive (default: 50) -- `--lr`: Learning rate (default: 0.005) -- `--margin`: Ranking loss margin (default: 0.2) -- `--lambda-reg`: Regularization strength (default: 0.1) -- `--gpu`: GPU device or -1 for CPU (default: -1) - -## Output Files - -After training, you'll get: -- `taxonomy_model_small_best.pth` - Best model (lowest loss) -- `taxonomy_model_small_epoch{N}.pth` - Last 5 epoch checkpoints -- `taxonomy_model_small.pth` - Final model - -## Prerequisites - -Make sure you've built the transitive closure first: -```bash -python build_transitive_closure.py -``` - -This creates `data/taxonomy_edges_small_transitive.pkl` which the training script needs. - -## Next Steps - -After training: - -1. **Analyze hierarchy quality:** - ```bash - python analyze_hierarchy_hyperbolic.py - ``` - -2. **Visualize embeddings:** - ```bash - python scripts/visualize_embeddings.py taxonomy_model_small_best.pth --highlight mammals - ``` - -## Troubleshooting - -### "Training data not found" -Run `python build_transitive_closure.py` first. - -### Training is too slow -- Reduce batch size: `--batch-size 32` -- Reduce negative samples: `--n-negatives 20` -- Use fewer epochs: `--epochs 50` - -### Embeddings escape the ball (Outside > 0%) -- Increase regularization: `--lambda-reg 0.2` -- Decrease learning rate: `--lr 0.001` - -### Loss not improving -- Increase learning rate: `--lr 0.01` -- Increase margin: `--margin 0.3` -- Train longer: `--epochs 200 --early-stopping 10` diff --git a/docs/archive/BALL_CONSTRAINT_ENFORCEMENT.md b/docs/archive/BALL_CONSTRAINT_ENFORCEMENT.md deleted file mode 100644 index 94fc28c..0000000 --- a/docs/archive/BALL_CONSTRAINT_ENFORCEMENT.md +++ /dev/null @@ -1,196 +0,0 @@ -# Ball Constraint Enforcement - 3-Layer Strategy - -## Problem Analysis (v2) - -After 1 epoch with improved settings: -``` -Total embeddings: 92,290 -Embeddings OUTSIDE ball (>=1.0): 2,051 (2.22%) -Embeddings INSIDE ball (<1.0): 90,239 (97.78%) - -Percentiles: - 95%: 0.9986 ✅ - 99%: 1.0356 ❌ (top 1% escaping!) - Max: 1.4522 ❌ -``` - -**Diagnosis:** Only the top 1-2% of embeddings escape, but they escape significantly (up to 1.45). - ---- - -## Solution: 3-Layer Enforcement (v3) - -### **Layer 1: Improved Projection (Per Batch)** -```python -# OLD: Soft clamp (affects all embeddings) -scale = torch.clamp(norms, max=1 - eps) / (norms + eps) - -# NEW: Hard constraint (only affects violators) -needs_projection = norms >= (1 - eps) -scale = torch.where( - needs_projection, - (1 - eps) / (norms + eps), - torch.ones_like(norms) -) -``` - -**Impact:** -- Only scales embeddings that are actually outside -- Leaves good embeddings untouched -- More precise enforcement - -### **Layer 2: Periodic Full Projection (Every 500 Batches)** -```python -if n_batches % 500 == 0: - model.project_to_ball(indices=None) # Project ALL embeddings -``` - -**Impact:** -- Catches stragglers that weren't in recent batches -- Happens ~30 times per epoch (15,248 batches / 500) -- Minimal overhead (~0.1% slowdown) - -### **Layer 3: Epoch-End Full Projection (Every Epoch)** -```python -# At end of each epoch, before saving checkpoint -model.project_to_ball(indices=None) - -# Verify and report -outside_count = (norms >= 1.0).sum().item() -if outside_count > 0: - print(f" ⚠️ {outside_count} embeddings still outside ball (should be 0!)") -``` - -**Impact:** -- **GUARANTEES** all saved checkpoints have valid embeddings -- Clear visibility if constraint is being violated -- Safety net before analysis - ---- - -## Complete Enforcement Stack - -| Layer | Trigger | Target | Purpose | -|-------|---------|--------|---------| -| **Gradient Clip** | Every batch | Gradients | Prevent exploding updates | -| **Regularizer** | Every batch | Loss | Soft penalty (encourages correct radii) | -| **Batch Projection** | Every batch | Updated embeddings | Hard constraint (immediate fix) | -| **Periodic Projection** | Every 500 batches | ALL embeddings | Catch stragglers | -| **Epoch Projection** | Every epoch | ALL embeddings | Guarantee checkpoint validity | - ---- - -## Expected Results (v3) - -### **After Epoch 1:** -``` -✅ Max norm: < 1.0 (was 1.45) -✅ Outside count: 0 (was 2,051) -✅ 100% embeddings inside ball (was 97.78%) -``` - -### **Training Characteristics:** - -1. **Stability:** ✅ - - No exploding gradients (clipped at 1.0) - - No runaway embeddings (3 projection layers) - -2. **Efficiency:** ✅ - - Most projections are per-batch (fast, only updated nodes) - - Full projections are rare (every 500 batches + epoch end) - - Overhead: ~2% slower than no projection - -3. **Correctness:** ✅ - - All saved checkpoints are valid - - Hyperbolic distances are well-defined - - No undefined/NaN values - ---- - -## Trade-offs - -### **Pros:** -- ✅ **Hard constraint:** Embeddings CANNOT escape -- ✅ **Multi-layer:** Redundant enforcement (belt + suspenders) -- ✅ **Verified:** Reports violations if they occur -- ✅ **Efficient:** Minimal overhead - -### **Cons:** -- ⚠️ **Optimization conflict:** Projection fights ranking loss - - Ranking loss wants some embeddings far apart - - Projection forces them back - - May slow convergence slightly - -- ⚠️ **Local minima risk:** Constrained optimization is harder - - Model has less freedom to explore - - May get stuck in suboptimal solution - -### **Mitigation:** -- Use soft regularizer (λ=0.1) + hard projection -- Regularizer guides gradients toward valid solutions -- Projection is safety net, not primary mechanism -- Learning rate (0.005) allows careful exploration - ---- - -## Comparison - -| Version | Max Norm | Outside Count | Strategy | -|---------|----------|---------------|----------| -| **v1 (Broken)** | 2.18 | ~50,000 (54%) | Weak reg (0.01), no grad clip | -| **v2 (Better)** | 1.45 | 2,051 (2.2%) | Strong reg (0.1), grad clip | -| **v3 (Strict)** | <1.0 | 0 (0%) | All of above + 3-layer projection | - ---- - -## Monitoring - -Watch for these patterns during training: - -### **Good Signs:** -- ✅ Max norm stays < 1.0 consistently -- ✅ Outside count = 0 every epoch -- ✅ Mean norm increases gradually (depth differentiation) -- ✅ Loss decreases steadily - -### **Warning Signs:** -- ⚠️ Max norm = 0.99999 many epochs (hitting boundary too hard) -- ⚠️ Loss plateaus early (over-constrained) -- ⚠️ Mean norm stays low (not learning depth structure) - -### **Fix if Over-Constrained:** -1. Reduce regularization: λ=0.05 (from 0.1) -2. Increase learning rate: 0.01 (from 0.005) -3. Remove periodic projection (keep only batch + epoch) - ---- - -## Code Changes (v2 → v3) - -### **File: train_hierarchical.py** - -**1. Improved projection logic (lines 104-135)** -- Changed from soft clamp to hard constraint -- Only scales embeddings that violate constraint - -**2. Periodic full projection (lines 379-382)** -- Added every 500 batches -- Projects ALL embeddings - -**3. Epoch-end full projection (lines 399-409)** -- Added before checkpoint save -- Reports violations - ---- - -## Summary - -**v3 uses a 3-layer defense strategy:** - -1. **Prevent:** Gradient clipping + strong regularization -2. **Correct:** Per-batch projection of updated embeddings -3. **Enforce:** Periodic + epoch-end full projection - -**Result:** ZERO embeddings outside ball, guaranteed valid checkpoints, mathematically sound Poincaré embeddings. - -This is the **proper way** to enforce manifold constraints in optimization! diff --git a/docs/archive/BUGS_FOUND_AND_FIXED.md b/docs/archive/BUGS_FOUND_AND_FIXED.md deleted file mode 100644 index 153fc16..0000000 --- a/docs/archive/BUGS_FOUND_AND_FIXED.md +++ /dev/null @@ -1,186 +0,0 @@ -# Critical Bugs Found and Fixed - -## Session: Nov 8, 2025 - Pre-Training Sanity Check - ---- - -## 🐛 BUG #1: TaxID vs Index Mapping Confusion (CRITICAL) - -### **Symptom:** -- Training created 3,467,244 embeddings instead of 111,103 -- Max index in data: 3,467,243 -- Expected max index: 111,102 -- Result: Training on mostly zero-gradient embeddings (wasted 97% of memory and compute) - -### **Root Cause:** -```python -# WRONG - build_transitive_closure.py line 173 -df = pd.read_csv("data/taxonomy_edges_small.mapping.tsv", - sep="\t", header=None, names=["idx", "taxid"]) -``` - -File has columns: `taxid idx` (with header) -Code expected: `idx taxid` (no header) - -Result: Code read TaxID 131567 as index 1 → used TaxIDs directly as indices! - -### **Fix:** -```python -# CORRECT -df = pd.read_csv("data/taxonomy_edges_small.mapping.tsv", sep="\t") -# Reads header automatically, columns are taxid, idx -``` - -### **Impact:** -- **Before:** 3.4M embeddings, 97% never updated -- **After:** 92K embeddings (actual max index from data) -- **Speedup:** ~30x less memory, proper training - ---- - -## ⚠️ WARNING #1: Index Range Mismatch - -### **Observation:** -- Mapping file has 111,103 entries (indices 0-111,102) -- Training data only uses indices 0-92,289 -- **18,814 mapped nodes (17%) never appear in training!** - -### **Cause:** -Some organisms in the small dataset have no valid taxonomy paths to root in the transitive closure. - -### **Status:** -Not a bug - these are likely: -- Organisms with missing parent links -- Isolated subgraphs -- Data quality issues in NCBI taxonomy - -### **Impact:** -- Minimal - these nodes won't have good embeddings -- Could improve by filling in missing links - ---- - -## ⚠️ WARNING #2: Regularization Coverage - -### **Observation:** -- 92,290 nodes in training data -- Only 82,040 nodes have depth info (89%) -- **10,250 nodes (11%) won't be regularized** - -### **Cause:** -Some nodes appear in training pairs but don't have depth metadata assigned. - -### **Status:** -Acceptable - these are likely intermediate nodes. Main concern is leaf nodes. - ---- - -## ✅ ALL OTHER SYSTEMS VERIFIED - -### **1. Mapping File** ✅ -- ✅ Columns: `taxid`, `idx` (correct order) -- ✅ No duplicate TaxIDs -- ✅ No duplicate indices -- ✅ Indices continuous: 0-111,102 - -### **2. Transitive Closure Data** ✅ -- ✅ 975,896 ancestor-descendant pairs -- ✅ All indices in valid range -- ✅ No self-loops -- ✅ Depth differences consistent -- ✅ All depth diffs positive - -### **3. Projection Logic** ✅ -- ✅ Correctly constrains embeddings to unit ball -- ✅ Max norm after projection: 0.99999 -- ✅ No embeddings escape ball - -### **4. Hyperbolic Distance** ✅ -- ✅ Distance to self ≈ 0 -- ✅ Distance increases with separation -- ✅ Proper Poincaré distance formula - -### **5. Depth-Aware Initialization** ✅ -- ✅ Root (depth 0): r = 0.10 -- ✅ Leaves (depth 38): r = 0.95 -- ✅ All radii < 1.0 (inside ball) - -### **6. Sibling Map (Hard Negatives)** ✅ -- ✅ 82,039 nodes with sibling info -- ✅ Average 31,438 siblings per node -- ✅ Only 1 node has no siblings -- ✅ Siblings verified at same depth - -### **7. Regularizer Targets** ✅ -- ✅ All targets < 1.0 (valid) -- ✅ Proper depth → radius mapping - -### **8. Training Configuration** ✅ -- ✅ 975,896 training pairs -- ✅ 15,249 batches per epoch -- ✅ Reasonable batch size (64) - ---- - -## Summary - -### **Critical Bugs Fixed:** -1. ✅ **TaxID as index bug** - Fixed mapping file reading - -### **Warnings (Non-Critical):** -1. ⚠️ 17% of mapped nodes never appear in training (data quality) -2. ⚠️ 11% of training nodes lack depth info (acceptable) - -### **Verification Results:** -- ✅ **10/10 sanity checks passed** -- ✅ **All core systems working correctly** -- ✅ **Ready to train with confidence** - ---- - -## Next Steps - -1. ✅ Run sanity check: `python sanity_check.py` -2. ⏭️ Train model: `python train_hierarchical.py ...` -3. ⏭️ Analyze results: `python analyze_hierarchy_hyperbolic.py` -4. ⏭️ Expect MUCH better results: - - Depth-norm correlation: r > 0.5 (was -0.08) - - Phylum separation: > 1.5x (was 1.05x) - - Class separation: > 1.5x (was 1.04x) - ---- - -## Files Modified - -1. **build_transitive_closure.py** - - Fixed mapping file reading (line 173-175) - - Now correctly reads header and column names - -2. **sanity_check.py** (NEW) - - Comprehensive validation of entire pipeline - - 8 test categories, 10 checks - - Run before every major training session - ---- - -## Lessons Learned - -1. **Always validate data format assumptions** - - CSV headers matter! - - Column order matters! - - Don't assume - verify! - -2. **Sanity check EVERYTHING before long training runs** - - Indices in range - - No self-loops - - Math checks out - - Data makes sense - -3. **Monitor actual vs expected sizes** - - 3.4M embeddings was a red flag - - Should have caught it earlier - -4. **Test mathematical operations independently** - - Projection logic - - Distance functions - - Initialization ranges diff --git a/docs/archive/CHECKPOINT_MANAGEMENT.md b/docs/archive/CHECKPOINT_MANAGEMENT.md deleted file mode 100644 index b072cdf..0000000 --- a/docs/archive/CHECKPOINT_MANAGEMENT.md +++ /dev/null @@ -1,144 +0,0 @@ -# Checkpoint Management - -## Overview - -To prevent disk space issues during training, the repository includes automatic checkpoint management that keeps only the **20 most recent checkpoints** per model. - -## Automatic Management During Training - -The `train_with_early_stopping.py` script includes built-in checkpoint management: - -```bash -python train_with_early_stopping.py -``` - -**Features:** -- ✅ Monitors training in real-time -- ✅ Implements early stopping (patience=6 epochs) -- ✅ Automatically keeps only 20 most recent checkpoints -- ✅ Deletes old checkpoints as new ones are created -- ✅ Shows space-saving notifications - -**Example Output:** -``` -💾 Checkpoint saved: taxonomy_model_small_early_stop_epoch12.pth (keeping 12/20) -💾 Checkpoint saved: taxonomy_model_small_early_stop_epoch13.pth (keeping 13/20) -... -💾 Checkpoint saved: taxonomy_model_small_early_stop_epoch21.pth (keeping 20/20) -🗑️ Deleted old checkpoint: taxonomy_model_small_early_stop_epoch1.pth -💾 Checkpoint saved: taxonomy_model_small_early_stop_epoch22.pth (keeping 20/20) -🗑️ Deleted old checkpoint: taxonomy_model_small_early_stop_epoch2.pth -``` - -## Manual Cleanup - -If you have old checkpoints from previous training runs, use the cleanup script: - -```bash -# See what would be deleted (dry run) -python cleanup_old_checkpoints.py --dry-run - -# Delete old checkpoints, keep 20 most recent -python cleanup_old_checkpoints.py - -# Keep only 10 most recent -python cleanup_old_checkpoints.py --keep 10 - -# Clean specific directory -python cleanup_old_checkpoints.py --directory /path/to/checkpoints -``` - -## Checkpoint Naming Formats - -The system recognizes these checkpoint patterns: - -1. **Epoch-based:** `model_name_epoch123.pth` -2. **Dot-notation:** `model_name.pth.123` - -Both formats are automatically detected and managed. - -## How It Works - -### During Training -1. Training saves checkpoints each epoch: `model_epoch0.pth`, `model_epoch1.pth`, etc. -2. The monitoring script tracks all saved checkpoints in a queue (max 20) -3. When a 21st checkpoint is saved, the oldest (epoch 0) is automatically deleted -4. This continues throughout training, maintaining exactly 20 checkpoints - -### Manual Cleanup -1. Script scans directory for checkpoint files -2. Groups checkpoints by model name -3. Sorts by epoch number -4. Keeps N most recent, deletes the rest - -## Disk Space Savings - -**Example for small dataset:** -- Each checkpoint: ~50 MB -- Without management: 200 epochs × 50 MB = **10 GB** -- With management: 20 checkpoints × 50 MB = **1 GB** -- **Savings: 9 GB (90%)** - -**Example for full dataset:** -- Each checkpoint: ~219 MB -- Without management: 200 epochs × 219 MB = **43.8 GB** -- With management: 20 checkpoints × 219 MB = **4.4 GB** -- **Savings: 39.4 GB (90%)** - -## Configuration - -### Change Maximum Checkpoints - -Edit `train_with_early_stopping.py`: - -```python -# Keep 30 checkpoints instead of 20 -summary = train_with_monitoring(cmd, patience=6, max_checkpoints=30, ...) -``` - -### Disable Auto-Cleanup - -If you want to keep all checkpoints, set a very high limit: - -```python -summary = train_with_monitoring(cmd, patience=6, max_checkpoints=9999, ...) -``` - -## Best Practices - -1. **Use early stopping:** Training typically converges in 20-50 epochs, so keeping 20 checkpoints is sufficient -2. **Run cleanup before training:** Clear out old experiments to start fresh -3. **Monitor disk space:** Use `df -h` to check available space -4. **Keep final model:** The script always saves the final model separately at the end - -## Files Involved - -- `train_with_early_stopping.py` - Training with auto-cleanup -- `cleanup_old_checkpoints.py` - Manual cleanup utility -- `hype/train.py` - Core training loop (saves per-epoch checkpoints) - -## Troubleshooting - -### Checkpoints Not Being Deleted - -Check if checkpoint pattern matches. The script looks for: -- `Saved checkpoint: ` in training output - -### Running Out of Disk Space - -1. Stop training (Ctrl+C) -2. Run manual cleanup: `python cleanup_old_checkpoints.py --keep 5` -3. Check space: `df -h` -4. Resume training - -### Accidentally Deleted Important Checkpoint - -The final model is always saved separately as `taxonomy_model_small_early_stop.pth` (without epoch number), which is never deleted. - -## Summary - -✅ **Automatic:** Training script manages checkpoints automatically -✅ **Efficient:** Saves 90% disk space -✅ **Safe:** Always keeps the 20 most recent checkpoints -✅ **Manual option:** Cleanup script for existing files -✅ **Configurable:** Adjust limits as needed diff --git a/docs/archive/CLEANUP_SUMMARY.md b/docs/archive/CLEANUP_SUMMARY.md deleted file mode 100644 index 52ceb38..0000000 --- a/docs/archive/CLEANUP_SUMMARY.md +++ /dev/null @@ -1,272 +0,0 @@ -# Repository Cleanup Summary - -## What Was Removed - -### Checkpoint Files -- **Removed:** 569 checkpoint files -- **File types:** `*.pth`, `*.pth.*` -- **Total size freed:** ~100+ GB - -### Log Files -- **Removed:** All log files -- **Files:** `training.log`, `training_full.log`, `nohup.out` - -### Visualization Files -- **Removed:** All generated PNG files -- **Files:** `umap_*.png`, `umap_projection.png`, etc. - -### Redundant Scripts -**Consolidated into `scripts/visualize_embeddings.py`:** -- ❌ `visualize_primates.py` -- ❌ `visualize_primates_proper.py` -- ❌ `visualize_primates_small_only.py` -- ❌ `visualize_by_taxonomy.py` -- ❌ `visualize_trained_small_dataset.py` - -**Old shell scripts removed:** -- ❌ `train-mammals.sh` -- ❌ `train-nouns.sh` -- ❌ `train_taxonomy.sh` -- ❌ `train_taxonomy_quick.sh` - -## New Universal Tools Created - -### 1. `scripts/visualize_embeddings.py` ⭐ -**Purpose:** One script to visualize any checkpoint - -**Features:** -- Works with any checkpoint file -- Highlight any taxonomic group (primates, mammals, bacteria, etc.) -- Show only specific groups -- Nearest neighbor analysis -- Automatic output naming -- Configurable sampling - -**Usage:** -```bash -# Basic -python scripts/visualize_embeddings.py model.pth - -# Highlight primates -python scripts/visualize_embeddings.py model.pth --highlight primates - -# Only show mammals -python scripts/visualize_embeddings.py model.pth --only mammals - -# Custom sample size -python scripts/visualize_embeddings.py model.pth --sample 50000 -``` - -### 2. `scripts/cleanup_repo.sh` -**Purpose:** Automated repository cleanup - -**Features:** -- Interactive confirmation -- Removes checkpoints, logs, visualizations -- Removes redundant scripts -- Reports what will be deleted - -**Usage:** -```bash -./scripts/cleanup_repo.sh -``` - -### 3. `scripts/validate_data.py` -**Purpose:** Data quality validation - -**Features:** -- Validates edgelist format -- Checks mapping consistency -- Verifies sequential indices -- Detects header bugs - -**Usage:** -```bash -python scripts/validate_data.py small -python scripts/validate_data.py full -``` - -## Repository Structure (After Cleanup) - -``` -taxembed/ -├── Core Scripts (Root) -│ ├── embed.py # Main training -│ ├── prepare_taxonomy_data.py # Data preparation -│ ├── remap_edges.py # Data remapping -│ ├── monitor_training.py # Training monitor -│ ├── evaluate_full.py # Evaluation -│ ├── evaluate_and_visualize.py # Combined eval -│ ├── nn_demo.py # Quick demo -│ └── reconstruction.py # Reconstruction eval -│ -├── src/taxembed/ # Source code -│ ├── manifolds/ # Hyperbolic manifolds -│ ├── models/ # Embedding models -│ ├── datasets/ # Data loading -│ └── utils/ # Utilities -│ -├── scripts/ # Organized utilities -│ ├── visualize_embeddings.py # ⭐ Universal visualization -│ ├── validate_data.py # ⭐ Data validation -│ ├── cleanup_repo.sh # ⭐ Repository cleanup -│ ├── regenerate_data.sh # Data regeneration -│ ├── prepare_data.py # Wrappers -│ ├── remap_data.py -│ ├── monitor.py -│ ├── evaluate.py -│ └── train.py -│ -├── tests/ # Unit tests -│ ├── __init__.py -│ └── test_example.py -│ -├── hype/ # Original package (backward compat) -│ -├── Configuration -│ ├── pyproject.toml # Project config (uv) -│ ├── ruff.toml # Linter config -│ ├── Makefile # Convenience commands -│ ├── setup.py # C++ extensions -│ ├── requirements.txt # Legacy requirements -│ └── .gitignore # Git ignore rules -│ -└── Documentation - ├── README.md # Main documentation - ├── QUICKSTART.md # Quick start guide - ├── GETTING_STARTED.md # Getting started - ├── STRUCTURE.md # Project structure - ├── SCRIPTS_GUIDE.md # ⭐ Script documentation - ├── CONTRIBUTING.md # Contribution guide - ├── DATA_FIXES_SUMMARY.md # Data bug fixes - ├── DATA_HANDLING_REVIEW.md # Data analysis - ├── CLEANUP_SUMMARY.md # This file - ├── RESTRUCTURING_SUMMARY.md # Restructuring notes - ├── RESTRUCTURING_COMPLETE.md # Restructuring completion - ├── PROJECT_TREE.txt # Visual tree - ├── TRAINING_SUMMARY.md # Training notes - ├── IMPLEMENTATION_NOTES.md # Implementation notes - ├── FINAL_ASSESSMENT.md # Quality assessment - └── CODE_OF_CONDUCT.md # Code of conduct -``` - -## Updated .gitignore - -Now properly ignores: -- Checkpoints: `*.pth`, `*.pth.*` -- Logs: `*.log`, `training*.log`, `nohup.out` -- Visualizations: `*.png`, `*.jpg` -- Data: `data/` -- Build artifacts: `build/`, `dist/`, `*.so` -- Python cache: `__pycache__/`, `*.pyc` -- Virtual environments: `venv/`, `venv311/` -- IDE files: `.idea/`, `.vscode/` - -## Benefits of Cleanup - -### Before Cleanup -- 569 checkpoint files (~100+ GB) -- 8 redundant visualization scripts -- 4 old shell scripts -- Numerous log and PNG files -- Confusing script organization - -### After Cleanup -- ✅ Clean repository -- ✅ 1 universal visualization tool (replaces 5 scripts) -- ✅ Clear script organization -- ✅ Comprehensive documentation -- ✅ Proper .gitignore -- ✅ Easy to maintain - -## Workflow Examples - -### Training a Model -```bash -python embed.py \ - -dset data/taxonomy_edges_small.mapped.edgelist \ - -checkpoint my_model.pth \ - -dim 10 -epochs 50 -negs 50 -burnin 10 \ - -batchsize 32 -model distance -manifold poincare \ - -lr 0.1 -gpu -1 -ndproc 1 -train_threads 1 \ - -eval_each 999999 -fresh -``` - -### Visualizing Results -```bash -# Highlight primates -python scripts/visualize_embeddings.py my_model.pth --highlight primates - -# Only show mammals -python scripts/visualize_embeddings.py my_model.pth --only mammals --sample 30000 - -# Basic visualization with nearest neighbors -python scripts/visualize_embeddings.py my_model.pth --nearest 10 -``` - -### Validating Data -```bash -python scripts/validate_data.py small -``` - -### Cleaning Up -```bash -./scripts/cleanup_repo.sh -``` - -## Key Improvements - -### 1. Consolidation -- **Before:** 5 separate visualization scripts, each hardcoded for specific use cases -- **After:** 1 universal tool that works with any checkpoint and any taxonomic group - -### 2. Documentation -- **Before:** Minimal script documentation -- **After:** Comprehensive `SCRIPTS_GUIDE.md` with usage examples - -### 3. Organization -- **Before:** Scripts scattered in root directory -- **After:** Organized in `scripts/` directory with clear purposes - -### 4. Maintenance -- **Before:** Hard to understand which scripts to use -- **After:** Clear documentation and single universal tool - -### 5. Disk Space -- **Before:** 100+ GB of old checkpoints -- **After:** Clean repository, generate files as needed - -## Future Maintenance - -### When Training -1. Train model: `python embed.py ...` -2. Visualize: `python scripts/visualize_embeddings.py --highlight ` -3. Clean up: `./scripts/cleanup_repo.sh` (when done) - -### When Adding Features -- Add to `scripts/` directory -- Update `SCRIPTS_GUIDE.md` -- Follow naming convention: `_.py` - -### When Sharing Code -- Repository is now clean and presentable -- Clear documentation for users -- No large binary files -- Professional organization - -## Recommendations - -1. **Use the universal visualization tool** for all embedding visualizations -2. **Clean up regularly** with `./scripts/cleanup_repo.sh` -3. **Validate data** before training with `scripts/validate_data.py` -4. **Follow the scripts guide** for standard workflows -5. **Keep documentation updated** when adding new scripts - -## Summary - -✅ **Removed:** 569 checkpoints, 8 redundant scripts, numerous temp files -✅ **Created:** Universal visualization tool, cleanup script, comprehensive documentation -✅ **Organized:** Scripts in proper directories, clear naming, good documentation -✅ **Professional:** Clean repo ready for production use and sharing - -The repository is now **production-ready** with a clean, maintainable structure! 🎉 diff --git a/docs/archive/DATA_FIXES_SUMMARY.md b/docs/archive/DATA_FIXES_SUMMARY.md deleted file mode 100644 index b0a4648..0000000 --- a/docs/archive/DATA_FIXES_SUMMARY.md +++ /dev/null @@ -1,188 +0,0 @@ -# Data Handling Fixes - Implementation Summary - -## Issues Found & Fixed - -### ❌ Bug 1: Header Lines Treated as Data -**Problem:** The `.edgelist` files had header lines ("id1 id2") that were being treated as actual organism TaxIDs. - -**Impact:** -- Old dataset: 111,105 nodes (included "id1" and "id2" as fake organisms) -- New dataset: 111,103 nodes (real organisms only) - -**Fix:** -- Updated `remap_edges.py` to detect and skip header lines -- Added validation for non-numeric values -- Headers like "id1", "id2", "taxid", and lines starting with "#" are now skipped - -### ⚠️ Bug 2: Inconsistent Mapping -**Problem:** The `.mapping.tsv` file contained the header strings as if they were TaxIDs. - -**Impact:** -``` -Old mapping: - taxid idx - id1 0 ← Fake TaxID - id2 1 ← Fake TaxID - 2 2 ← First real organism - -New mapping: - taxid idx - 2 0 ← First real organism starts at index 0 - 131567 1 ← Correct sequential mapping -``` - -**Fix:** -- Fixed `remap_edges.py` automatically handles this -- Mapping file now only contains real TaxIDs - -### ✅ Improvement: Better Error Handling -**Added:** -- Input validation for numeric values -- Better error messages -- Progress reporting during remapping - -## Files Modified - -### 1. `remap_edges.py` -**Changes:** -- Added header detection and skipping -- Added numeric validation -- Better error messages -- Progress output -- Made into proper Python module with `main()` - -### 2. `prepare_taxonomy_data.py` -**Changes:** -- Now creates `.edgelist` files directly (without headers) -- Creates both CSV (with header) and edgelist (without header) formats - -### 3. New: `scripts/validate_data.py` -**Purpose:** -- Validates edgelist files for headers and numeric values -- Validates mapping files for consistency -- Checks node indices are sequential -- Verifies consistency between edgelist and mapping - -**Usage:** -```bash -python scripts/validate_data.py small # Validate small dataset -python scripts/validate_data.py full # Validate full dataset -``` - -### 4. New: `scripts/regenerate_data.sh` -**Purpose:** -- Automated script to regenerate all data files -- Runs validation automatically -- Creates both full and small datasets - -## Validation Results - -### Before Fixes (Buggy Data) -``` -Small dataset validation: - ❌ Found 2 non-numeric TaxIDs: 'id1', 'id2' - ⚠️ 2 nodes in edgelist not in mapping -``` - -### After Fixes (Clean Data) -``` -Small dataset: - ✅ 111,103 nodes (removed 2 fake nodes) - ✅ 100,000 edges - ✅ All checks passed - -Full dataset: - ✅ 2,705,745 nodes - ✅ 2,705,744 edges - ✅ All checks passed -``` - -## Impact on Training - -### Before (with bugs): -- Training on 111,105 nodes (2 fake + 111,103 real) -- Indices 0 and 1 were "id1" and "id2" strings -- Actual organisms started at index 2 - -### After (clean): -- Training on 111,103 real organisms -- Indices 0-111,102 are all real TaxIDs -- No fake nodes affecting embeddings - -### Model Performance: -**The bugs had minimal impact** on the trained model because: -- The 2 fake nodes had very few edges -- They didn't participate meaningfully in training -- The model still learned hierarchical structure correctly - -**However, the fixes provide:** -- ✅ Cleaner data -- ✅ Correct node counts -- ✅ Better interpretability -- ✅ Easier debugging -- ✅ Reproducibility - -## Recommendations - -### For Future Training - -1. **Always validate data before training:** -```bash -python scripts/validate_data.py small -python scripts/validate_data.py full -``` - -2. **Retrain models on clean data:** - - The old models work fine, but for publication/production use clean data - - Expected slight difference in node count but same quality - -3. **Use validation in CI/CD:** - - Add data validation to your testing pipeline - - Prevents data quality regressions - -## Files Generated - -### Clean Data Files -- ✅ `data/taxonomy_edges_small.mapped.edgelist` (100,000 edges, 111,103 nodes) -- ✅ `data/taxonomy_edges_small.mapping.tsv` (111,103 mappings) -- ✅ `data/taxonomy_edges.mapped.edgelist` (2.7M edges, 2.7M nodes) -- ✅ `data/taxonomy_edges.mapping.tsv` (2.7M mappings) - -### Validation Scripts -- ✅ `scripts/validate_data.py` - Validates data quality -- ✅ `scripts/regenerate_data.sh` - Regenerates all data - -### Documentation -- ✅ `DATA_HANDLING_REVIEW.md` - Detailed analysis of issues -- ✅ `DATA_FIXES_SUMMARY.md` - This file - -## Testing - -To verify everything works: - -```bash -# 1. Validate data -python scripts/validate_data.py small - -# 2. Train a quick test (5 epochs) -python embed.py \ - -dset data/taxonomy_edges_small.mapped.edgelist \ - -checkpoint test_clean.pth \ - -dim 10 -epochs 5 -negs 50 -burnin 2 \ - -batchsize 32 -model distance -manifold poincare \ - -lr 0.1 -gpu -1 -ndproc 1 -train_threads 1 \ - -eval_each 999999 -fresh - -# 3. Check node count in checkpoint -python -c "import torch; c=torch.load('test_clean.pth'); print(f'Nodes: {len(c[\"objects\"]):,}')" -# Should show: Nodes: 111,103 (not 111,105) -``` - -## Conclusion - -✅ **All data handling bugs have been fixed** -✅ **Clean datasets regenerated and validated** -✅ **Validation tools created for future use** -✅ **Documentation updated** - -The repository now has **production-quality data handling** with proper validation and error checking. diff --git a/docs/archive/DATA_HANDLING_REVIEW.md b/docs/archive/DATA_HANDLING_REVIEW.md deleted file mode 100644 index 82c27be..0000000 --- a/docs/archive/DATA_HANDLING_REVIEW.md +++ /dev/null @@ -1,243 +0,0 @@ -# Data Handling Review - -## Current Pipeline - -### 1. Data Preparation (`prepare_taxonomy_data.py`) -``` -Input: data/nodes.dmp, data/names.dmp -Output: data/taxonomy_edges.csv -Format: CSV with header (id1,id2,weight) -``` - -**✅ This step is correct:** -- Extracts parent-child relationships from NCBI taxonomy -- Each edge: child → parent -- Skips self-loops (root node) -- Saves as CSV with proper header - -### 2. CSV → EdgeList Conversion -``` -Input: data/taxonomy_edges.csv -Output: data/taxonomy_edges.edgelist -Format: Whitespace-separated, WITH HEADER -``` - -**⚠️ ISSUE FOUND:** -```bash -$ head -5 data/taxonomy_edges_small.edgelist -id1 id2 # ← HEADER LINE (should be removed!) -2 131567 -6 335928 -7 6 -9 32199 -``` - -The `.edgelist` file contains a header line "id1 id2" which should NOT be there. - -### 3. Remapping (`remap_edges.py`) -```python -in_path = sys.argv[1] -out_edges = in_path.replace(".edgelist", ".mapped.edgelist") -out_map = in_path.replace(".edgelist", ".mapping.tsv") - -with open(in_path) as f: - for ln, line in enumerate(f, 1): - line=line.strip() - if not line: continue - parts = line.split() - if len(parts)!=2: - raise ValueError(f"Bad line {ln}: {line!r}") - u, v = parts - if u not in nodes: nodes[u] = len(nodes) - if v not in nodes: nodes[v] = len(nodes) - edges.append((nodes[u], nodes[v])) -``` - -**❌ BUG: Treats header as data!** - -Result: -```bash -$ head -5 data/taxonomy_edges_small.mapping.tsv -taxid idx -id1 0 # ← HEADER treated as a TaxID! -id2 1 # ← HEADER treated as a TaxID! -2 2 # ← Actual TaxID -131567 3 # ← Actual TaxID -``` - -This means: -- Node 0 = "id1" (string literal, not a taxid!) -- Node 1 = "id2" (string literal, not a taxid!) -- Node 2 = TaxID 2 (actual organism) -- Node 3 = TaxID 131567 (actual organism) - -### 4. Training Data Loading (`embed.py` → `hype/graph.py`) - -```python -def load_edge_list(path, symmetrize=False): - df = pandas.read_csv(path, sep=r'\s+', header=None, names=['id1', 'id2'], engine='python') - df['weight'] = 1.0 - df.dropna(inplace=True) - idx, objects = pandas.factorize(df[['id1', 'id2']].values.reshape(-1)) - idx = idx.reshape(-1, 2).astype('int') - weights = df.weight.values.astype('float') - return idx, objects.tolist(), weights -``` - -**Issue: Double Remapping** -- The `.mapped.edgelist` already has sequential indices (0, 1, 2, ...) -- `pandas.factorize` RE-MAPS them again based on first appearance order -- This creates a NEW mapping that may differ from the original mapping file - -**Result:** -- The `objects` array returned is the remapped values (not the original TaxIDs) -- The mapping file (`taxonomy_edges_small.mapping.tsv`) becomes INVALID for lookup -- When we try to map back using the .mapping.tsv file, we get wrong organisms - -## Problems Identified - -### Problem 1: Header Line in EdgeList -**Severity: HIGH** -**Impact:** Two fake "organisms" (id1, id2) are added to the dataset - -**Fix:** -```python -# In remap_edges.py, skip header line: -with open(in_path) as f: - for ln, line in enumerate(f, 1): - line = line.strip() - if not line or line.startswith("id1"): # Skip header - continue - # ... rest of code -``` - -Or better: Don't include header in .edgelist file at all. - -### Problem 2: Double Remapping -**Severity: MEDIUM** -**Impact:** Redundant computation, potential index mismatch - -**Current Flow:** -``` -TaxIDs → remap_edges.py → [0,1,2,...,N] → load_edge_list → factorize → [0,1,2,...,M] -``` - -The `factorize` step creates a NEW mapping based on first appearance order in the edge list. - -**Why this mostly works:** -- If edges are processed in order, factorize will likely preserve the mapping -- But it's not guaranteed and adds unnecessary complexity - -**Better approach:** -```python -def load_edge_list_remapped(path, symmetrize=False): - """Load edge list that's already been remapped to sequential indices.""" - edges = [] - max_idx = -1 - - with open(path) as f: - for line in f: - line = line.strip() - if not line or line.startswith("id1"): # Skip empty/header - continue - parts = line.split() - if len(parts) >= 2: - u, v = int(parts[0]), int(parts[1]) - edges.append([u, v]) - max_idx = max(max_idx, u, v) - - if symmetrize: - edges.extend([[v, u] for u, v in edges]) - - idx = np.array(edges, dtype=int) - # Create objects list: [0, 1, 2, ..., max_idx] - objects = list(range(max_idx + 1)) - weights = np.ones(len(idx), dtype=float) - - return idx, objects, weights -``` - -### Problem 3: Mapping File Inconsistency -**Severity: HIGH** -**Impact:** Cannot correctly map embeddings back to TaxIDs - -The `.mapping.tsv` file maps: `TaxID → original_idx` -But the training uses: `factorized_idx → embedding` - -To get `TaxID → embedding`, you need: -1. `TaxID → original_idx` (from .mapping.tsv) -2. `original_idx → factorized_idx` (not saved anywhere!) -3. `factorized_idx → embedding` (from model) - -**Current workaround:** -The training saves `objects` in the checkpoint, which are the factorized values. So you can use: -- `checkpoint['objects'][factorized_idx]` → `original_idx_str` -- Then lookup in mapping file - -But this is confusing and error-prone. - -## Recommendations - -### Immediate Fixes - -1. **Remove header from .edgelist files:** -```bash -# Convert CSV to edgelist without header -tail -n +2 data/taxonomy_edges.csv | awk '{print $1, $2}' FS=',' > data/taxonomy_edges.edgelist -``` - -2. **Update remap_edges.py to skip headers:** -```python -if line.startswith("id1") or line.startswith("taxid"): - continue -``` - -3. **Verify data integrity:** -```bash -# Check that mapping file has no "id1" or "id2" -grep -E "^(id1|id2)\s" data/*.mapping.tsv -``` - -### Long-term Improvements - -1. **Simplify the pipeline:** - - Either use factorize OR pre-remap, not both - - Make load_edge_list aware that data is already remapped - -2. **Add data validation:** - - Check for header lines - - Verify node indices are sequential - - Confirm mapping file matches edge list - -3. **Improve documentation:** - - Document expected file formats - - Explain mapping strategy - - Add validation scripts - -## Current Status - -**Does it work?** YES, mostly. - -Despite the issues: -- The model trains successfully -- Loss decreases appropriately -- Embeddings learn hierarchical structure - -**Why it works despite bugs:** -- The header line adds 2 fake nodes, but they have few/no edges -- Double remapping is redundant but doesn't break functionality (if order preserved) -- The factorized `objects` are saved in checkpoints for reverse lookup - -**Should you fix it?** YES! -- Removes confusion -- Prevents potential bugs with different data -- Makes code more maintainable -- Improves reproducibility - -## Testing Recommendations - -After fixes: -1. Retrain small model (5-10 epochs) -2. Verify node count matches expectations -3. Check that embedding[0] corresponds to correct TaxID -4. Validate nearest neighbors match biological taxonomy diff --git a/docs/archive/FINAL_ASSESSMENT.md b/docs/archive/FINAL_ASSESSMENT.md deleted file mode 100644 index 039c579..0000000 --- a/docs/archive/FINAL_ASSESSMENT.md +++ /dev/null @@ -1,110 +0,0 @@ -# Final Assessment: Corners Cut vs. Production Quality - -## ✅ What We Did NOT Cut Corners On - -### 1. **Data Integrity** ✓ -- Used complete NCBI taxonomy dataset (2.7M organisms, 2.7M edges) -- Proper parent-child relationship preservation -- No data filtering or sampling during training -- Full hierarchical structure maintained - -### 2. **Model Training** ✓ -- 50 full epochs (standard for embeddings) -- Proper burn-in phase (10 epochs with reduced learning rate) -- Correct hyperparameters: - - Learning rate: 0.3 (appropriate for Poincaré) - - Batch size: 32 (reasonable for CPU) - - Negative samples: 50 (standard) - - Embedding dimension: 10 (as requested) -- Riemannian SGD optimizer (correct for hyperbolic space) -- Poincaré manifold (hyperbolic geometry preserved) - -### 3. **Code Quality** ✓ -- Fixed all critical bugs (elapsed time, device selection, format parsing) -- Proper error handling and fallbacks -- macOS-compatible (single-threaded, CPU-safe) -- Reproducible (fixed random seeds) - -## ⚠️ Where We DID Cut Corners (and Why) - -### 1. **Evaluation During Training** ✗ -- **What we cut**: Reconstruction evaluation (`-eval_each 999999`) -- **Why**: The evaluation code had indexing bugs that crashed on raw TaxIDs -- **Impact**: MINIMAL - evaluation is post-hoc, not required for training quality -- **Mitigation**: We ran evaluation AFTER training completed successfully -- **Status**: ✓ FIXED - evaluation works now - -### 2. **Multiprocessing** ✗ -- **What we cut**: Used single-threaded training (`-train_threads 1`, `-ndproc 0`) -- **Why**: Multiprocessing had serialization issues on macOS -- **Impact**: MINIMAL - training still completes in ~1 hour -- **Mitigation**: Single-threaded is actually more stable for development -- **Status**: ✓ ACCEPTABLE - can be re-enabled on Linux/proper cluster - -### 3. **GPU Acceleration** ✗ -- **What we cut**: Used CPU-only training -- **Why**: MPS (Apple Silicon GPU) had sparse operation incompatibilities -- **Impact**: MODERATE - ~10-20x slower than GPU, but still reasonable -- **Mitigation**: Can be re-enabled with `export PYTORCH_ENABLE_MPS_FALLBACK=1` -- **Status**: ✓ ACCEPTABLE - CPU training is stable and reproducible - -## 📊 Quality Metrics - -### Training Convergence -- ✓ No divergence observed -- ✓ Loss decreased smoothly across epochs -- ✓ Model saved successfully - -### Embedding Quality -- ✓ Nearest neighbors make biological sense - - Human → other primates - - Mouse → other rodents - - E. coli → other bacteria -- ✓ Hierarchical structure preserved -- ✓ UMAP visualization shows clear clustering - -### Reproducibility -- ✓ Fixed random seeds -- ✓ Deterministic on CPU -- ✓ All code changes documented -- ✓ Training log saved - -## 🎯 Production Readiness Assessment - -| Aspect | Status | Notes | -|--------|--------|-------| -| Data Completeness | ✅ | Full 2.7M organism dataset | -| Model Training | ✅ | 50 epochs, proper hyperparameters | -| Convergence | ✅ | Smooth, no divergence | -| Evaluation | ✅ | Post-training evaluation works | -| Reproducibility | ✅ | Fixed seeds, deterministic | -| Documentation | ✅ | Complete with examples | -| Visualization | ✅ | UMAP shows clear structure | -| **Overall** | **✅ PRODUCTION READY** | **Minor optimizations possible** | - -## 🚀 Recommended Next Steps (Optional Improvements) - -### High Priority -1. **Re-enable evaluation during training** - Now that bugs are fixed -2. **Test on GPU** - Use MPS with fallback enabled -3. **Multiprocessing** - Enable on Linux/cluster environments - -### Medium Priority -1. **Higher dimensions** - Train 50-100D embeddings for better representation -2. **Longer training** - 100+ epochs for convergence analysis -3. **Hyperparameter tuning** - Grid search for optimal LR, batch size - -### Low Priority -1. **Sparse gradients** - Enable for memory efficiency -2. **Symmetrization** - Train on bidirectional edges -3. **Fine-tuning** - Domain-specific adaptation - -## 📝 Summary - -**We trained a production-quality Poincaré embedding model on the complete NCBI taxonomy.** The corners we cut (evaluation during training, multiprocessing, GPU) were: -- Necessary for macOS compatibility -- Non-critical for model quality -- Easily reversible on proper infrastructure -- Well-documented and understood - -**The model is ready for downstream use** in taxonomic classification, species similarity search, and phylogenetic analysis. diff --git a/docs/archive/GETTING_STARTED.md b/docs/archive/GETTING_STARTED.md deleted file mode 100644 index fdc0dd5..0000000 --- a/docs/archive/GETTING_STARTED.md +++ /dev/null @@ -1,168 +0,0 @@ -# Getting Started with taxembed - -Welcome to the restructured taxembed project! This guide will help you get up and running. - -## 🚀 Quick Setup (5 minutes) - -### 1. Install Dependencies -```bash -make install -# This runs: uv sync -``` - -### 2. Build C++ Extensions -```bash -make build -# This runs: python setup.py build_ext --inplace -``` - -### 3. Verify Installation -```bash -make test -# This runs: uv run pytest -``` - -Done! You're ready to go. - -## 📖 What's New? - -### Project Structure -``` -taxembed/ -├── src/taxembed/ ← Main package code -├── scripts/ ← Executable scripts -├── tests/ ← Unit tests -├── pyproject.toml ← Dependencies (uv) -├── ruff.toml ← Code quality rules -└── Makefile ← Convenient commands -``` - -### Key Tools - -| Tool | Purpose | Command | -|------|---------|---------| -| **uv** | Fast package manager | `make install` | -| **ruff** | Linter & formatter | `make lint`, `make format` | -| **pytest** | Testing framework | `make test` | -| **Makefile** | Command shortcuts | `make help` | - -## 🎯 Common Tasks - -### Training a Model -```bash -uv run python scripts/train.py \ - --dataset data/taxonomy_edges.mapped.edgelist \ - --checkpoint model.pth \ - --epochs 50 -``` - -### Checking Code Quality -```bash -make lint # Check for issues -make format # Auto-fix issues -``` - -### Running Tests -```bash -make test # Run all tests -make test-cov # With coverage report -``` - -### Cleaning Up -```bash -make clean # Remove build artifacts -``` - -## 📚 Documentation - -- **README.md** - Full project documentation -- **QUICKSTART.md** - Detailed quick start guide -- **STRUCTURE.md** - Project organization -- **CONTRIBUTING.md** - Development guidelines -- **RESTRUCTURING_SUMMARY.md** - What changed and why - -## 🔧 Development Workflow - -### Step 1: Make Changes -Edit files in `src/taxembed/` or `scripts/` - -### Step 2: Check Code Quality -```bash -make lint # Find issues -make format # Fix automatically -``` - -### Step 3: Test Your Changes -```bash -make test # Run tests -``` - -### Step 4: Commit -```bash -git add . -git commit -m "Your message" -``` - -## ⚡ Useful Commands - -```bash -# Show all available commands -make help - -# Install dependencies -make install - -# Build C++ extensions -make build - -# Check code style -make lint - -# Fix code style -make format - -# Run tests with coverage -make test-cov - -# Clean build artifacts -make clean - -# Run a script -uv run python scripts/train.py --help -``` - -## 🐛 Troubleshooting - -### "Command not found: uv" -Install uv: https://github.com/astral-sh/uv#installation - -### "ModuleNotFoundError: No module named 'hype'" -Run: `make install && make build` - -### "ruff: command not found" -Use: `uv run ruff check src/` (with `uv run` prefix) - -### "Tests fail" -Check that dependencies are installed: `make install` - -## 📝 Next Steps - -1. Read [QUICKSTART.md](QUICKSTART.md) for detailed instructions -2. Check [STRUCTURE.md](STRUCTURE.md) to understand the layout -3. Review [CONTRIBUTING.md](CONTRIBUTING.md) for development guidelines -4. Start training: `uv run python scripts/train.py --help` - -## 💡 Tips - -- Use `make help` to see all available commands -- Use `uv run` to execute Python scripts in the project environment -- Use `make format` before committing to maintain code style -- Use `make test` to verify your changes work - -## 🤝 Need Help? - -- Check the documentation files (README.md, STRUCTURE.md, etc.) -- Review CONTRIBUTING.md for development guidelines -- Open an issue on GitHub - -Happy coding! 🎉 diff --git a/docs/archive/IMPLEMENTATION_NOTES.md b/docs/archive/IMPLEMENTATION_NOTES.md deleted file mode 100644 index c92ff13..0000000 --- a/docs/archive/IMPLEMENTATION_NOTES.md +++ /dev/null @@ -1,158 +0,0 @@ -# Poincaré Embeddings for NCBI Taxonomy - Implementation Notes - -## Overview -This project implements Poincaré embeddings to learn hierarchical representations of the NCBI taxonomy dataset. The goal is to embed organisms in hyperbolic space such that taxonomically related organisms are closer together. - -## Key Modifications Made - -### 1. Data Loading & Batch Processing -- **Fixed**: `BatchedDataset` requires `num_workers > 0` to start worker threads - - Previous: `-ndproc 0` disabled threading, resulting in empty batches - - Solution: Use `-ndproc 1` to enable single worker thread - - File: `hype/graph_dataset.pyx` - -### 2. Training Loop Improvements -- **Added**: Per-epoch checkpoint saving for early stopping - - Saves `taxonomy_model_full_fixed_epochN.pth` after each epoch - - Enables monitoring and early stopping without losing progress - - File: `hype/train.py` (lines 96-104) - -- **Fixed**: `UnboundLocalError` for `elapsed` variable - - Initialized before batch loop - - File: `hype/train.py` - -### 3. Embedding Initialization (CRITICAL FIX) -- **Problem**: Embeddings initialized with scale=1e-4 ([-1e-4, 1e-4]) - - Initial norm ≈ 3e-4 (essentially zero) - - Gradients too small to learn - - Loss stayed constant at 3.931 - -- **Solution**: Changed initialization scale to 0.1 - - Initial norm ≈ 0.316 (reasonable) - - Gradients large enough for learning - - File: `hype/manifolds/manifold.py` (line 26) - -### 4. Data Format Support -- **Added**: Support for `.edgelist` file format - - Auto-detects file format (CSV vs whitespace-separated) - - File: `embed.py` - -### 5. Model Checkpointing -- **Added**: Final model save when evaluation is disabled - - Prevents loss of trained model - - File: `embed.py` (lines 309-314) - -## Training Parameters - -### Recommended Settings -```bash -python embed.py \ - -dset data/taxonomy_edges.mapped.edgelist \ - -checkpoint taxonomy_model_full_fixed.pth \ - -dim 10 \ - -epochs 50 \ - -negs 50 \ - -burnin 10 \ - -batchsize 32 \ - -model distance \ - -manifold poincare \ - -lr 0.1 \ - -gpu -1 \ - -ndproc 1 \ - -train_threads 1 \ - -eval_each 999999 \ - -fresh -``` - -### Learning Rate Schedule -- **Burn-in (epochs 0-9)**: lr = 0.1 × 0.01 = 0.001 (stable initialization) -- **Normal (epochs 10+)**: lr = 0.1 (faster learning) - -## Monitoring Training - -### Real-time Clustering Quality -```bash -python monitor_training.py -``` - -Shows per-epoch: -- Loss (should decrease) -- Primate distance (within-group) -- Random distance (all organisms) -- Ratio (random/primate) - should be > 1.5 for good clustering - -### Training Progress -```bash -tail -f training_full.log -``` - -## Dataset - -### NCBI Taxonomy Data -- **Nodes**: 2,705,747 organisms -- **Edges**: 2,705,745 parent-child relationships -- **Format**: Whitespace-separated edgelist with TaxID remapping -- **Files**: - - `data/taxonomy_edges.mapped.edgelist`: Full dataset - - `data/taxonomy_edges.mapping.tsv`: TaxID ↔ index mapping - - `data/nodes.dmp`: NCBI taxonomy node information - -### Data Preparation -```bash -python remap_edges.py -python prepare_taxonomy_data.py -``` - -## Known Issues & Limitations - -1. **Training Speed**: CPU training is slow (~40 min per epoch on 2.7M edges) - - Consider using GPU with `-gpu 0` (requires CUDA) - - Or use smaller dataset for testing - -2. **Hierarchical Learning**: Model requires many epochs to learn taxonomy - - Burn-in phase helps but adds overhead - - Consider reducing epochs for faster iteration - -3. **Memory Usage**: Full model checkpoint is ~206MB - - Each epoch saves separate checkpoint - - Clean up old checkpoints if storage is limited - -## Files Modified - -- `embed.py`: Main training script -- `hype/train.py`: Training loop with per-epoch checkpointing -- `hype/manifolds/manifold.py`: Fixed embedding initialization -- `hype/graph_dataset.pyx`: Batch generation (no changes needed, but requires `ndproc > 0`) -- `hype/graph.py`: Data loading utilities -- `requirements.txt`: Dependencies - -## New Files Added - -- `monitor_training.py`: Real-time clustering quality monitoring -- `TRAINING_SUMMARY.md`: Training results and metrics -- `FINAL_ASSESSMENT.md`: Quality assessment -- `evaluate_full.py`: Model evaluation and UMAP visualization -- `visualize_primates_proper.py`: Primate-specific visualization -- `prepare_taxonomy_data.py`: Data preparation utilities -- `remap_edges.py`: TaxID remapping utility - -## Next Steps - -1. **GPU Training**: Use GPU for faster training - - Modify `-gpu 0` parameter - - Set `export PYTORCH_ENABLE_MPS_FALLBACK=1` for macOS - -2. **Hyperparameter Tuning**: - - Experiment with embedding dimension (currently 10) - - Adjust learning rate and burn-in multiplier - - Try different negative sampling counts - -3. **Evaluation**: - - Run `python evaluate_full.py` to get UMAP visualizations - - Check nearest neighbors for sample organisms - - Measure reconstruction quality - -4. **Production Deployment**: - - Save best epoch checkpoint - - Create inference script for embedding new organisms - - Benchmark performance on downstream tasks diff --git a/docs/archive/NEXT_STEPS.md b/docs/archive/NEXT_STEPS.md deleted file mode 100644 index 9539229..0000000 --- a/docs/archive/NEXT_STEPS.md +++ /dev/null @@ -1,414 +0,0 @@ -# Next Steps: Extensions & Enhancements - -## ✅ Current Status - -**Repository is production-ready:** -- ✅ Clean, restructured codebase -- ✅ Data bugs fixed (111,103 clean nodes) -- ✅ Training validated (500 epochs, excellent results) -- ✅ Universal visualization tool -- ✅ Comprehensive documentation -- ✅ Updated README with usage-first approach - -**Current Capabilities:** -- Train on NCBI taxonomy graph structure -- Learn hierarchical embeddings in hyperbolic space -- Visualize any taxonomic group -- Nearest neighbor queries -- Data validation - -## 🚀 Proposed Extensions - -### 1. Species Names Integration - -**Goal:** Enable text-based queries and better organism search - -**Current State:** -- Names stored in mapping file -- NOT used in training -- Only for visualization - -**Extension:** -```python -# Multi-modal training -class MultiModalPoincare: - def __init__(self): - self.graph_encoder = PoincareEmbedding() # Current - self.text_encoder = BertModel() # NEW - - def forward(self, taxid, name_text): - graph_emb = self.graph_encoder(taxid) - text_emb = self.text_encoder(name_text) - - # Align embeddings - loss = distance(graph_emb, text_emb) - return loss -``` - -**Benefits:** -- Query: "find species similar to 'sapiens'" -- Handle synonyms: "Escherichia coli" = "E. coli" -- Cross-lingual: search in any language -- Better disambiguation - -**Implementation Steps:** -1. Load species names from `names.dmp` -2. Add BERT/BioBERT text encoder -3. Create text dataset: `(taxid, name)` pairs -4. Joint training: graph loss + text alignment loss -5. Update visualization to support text queries - -**Effort:** Medium (1-2 weeks) - ---- - -### 2. Protein Embeddings - -**Goal:** Incorporate functional/sequence information - -**Current State:** -- Only taxonomy structure -- No molecular information - -**Extension:** -```python -# Add protein-level features -class ProteinEnhancedEmbedding: - def __init__(self): - self.taxonomy_encoder = PoincareEmbedding() # Current - self.protein_encoder = ESMModel() # NEW (protein language model) - self.fusion = FusionLayer() # Combine both - - def forward(self, taxid, protein_sequences): - tax_emb = self.taxonomy_encoder(taxid) - - # Aggregate proteins for organism - protein_embs = [self.protein_encoder(seq) for seq in protein_sequences] - org_protein_emb = aggregate(protein_embs) # Mean/max pooling - - # Combine - combined_emb = self.fusion(tax_emb, org_protein_emb) - return combined_emb -``` - -**Data Sources:** -- UniProt for protein sequences -- KEGG for protein functions -- RefSeq for reference proteomes - -**Benefits:** -- Find organisms by protein function -- Cluster by functional similarity -- Better for organisms with horizontal gene transfer -- Useful for drug discovery - -**Implementation Steps:** -1. Download protein sequences from UniProt -2. Compute ESM/ProtT5 embeddings for each protein -3. Aggregate per organism (mean pooling) -4. Create fusion architecture -5. Joint training on taxonomy + proteins - -**Effort:** High (3-4 weeks) - ---- - -### 3. Additional Features - -**Goal:** Incorporate phenotypic/genomic metadata - -**Current State:** -- Only taxonomy relationships -- No organism features - -**Extension:** -```python -# Add feature vectors -class FeatureAugmentedEmbedding: - def __init__(self): - self.taxonomy_encoder = PoincareEmbedding() # Current - self.feature_encoder = FeatureNet() # NEW - - def forward(self, taxid, features): - tax_emb = self.taxonomy_encoder(taxid) - - # Features: [genome_size, gc_content, temperature, ...] - feat_emb = self.feature_encoder(features) - - # Concatenate or fuse - combined = torch.cat([tax_emb, feat_emb], dim=-1) - return combined -``` - -**Possible Features:** -- **Genomic:** genome size, GC content, chromosome count -- **Environmental:** temperature range, pH range, habitat -- **Morphological:** cell shape, motility, gram staining -- **Metabolic:** aerobic/anaerobic, carbon source - -**Data Sources:** -- NCBI BioSample -- IMG (DOE Joint Genome Institute) -- BacDive (bacterial metadata) - -**Benefits:** -- Predict missing features -- Find organisms by phenotype -- Better for ecological studies -- Feature-based queries - -**Implementation Steps:** -1. Collect feature data from databases -2. Create feature vectors per organism -3. Add feature encoder network -4. Train with feature prediction as auxiliary task -5. Enable feature-based search - -**Effort:** Medium (2-3 weeks) - ---- - -### 4. Word Descriptions - -**Goal:** Natural language descriptions and semantic search - -**Current State:** -- No text descriptions -- No literature linkage - -**Extension:** -```python -# Add descriptions -class DescriptionEnhancedEmbedding: - def __init__(self): - self.taxonomy_encoder = PoincareEmbedding() # Current - self.description_encoder = SentenceBERT() # NEW - - def forward(self, taxid, description_text): - tax_emb = self.taxonomy_encoder(taxid) - - # Encode description: "Gram-negative bacterium found in..." - desc_emb = self.description_encoder(description_text) - - # Align - loss = distance(tax_emb, desc_emb) - return loss -``` - -**Data Sources:** -- Wikipedia organism articles -- NCBI organism descriptions -- Literature abstracts (PubMed) -- Textbooks and databases - -**Benefits:** -- Semantic search: "find bacteria that ferment lactose" -- Generate organism descriptions -- Link to scientific literature -- Educational applications - -**Implementation Steps:** -1. Scrape/download organism descriptions -2. Add Sentence-BERT encoder -3. Create description dataset -4. Joint training on taxonomy + descriptions -5. Enable natural language queries - -**Effort:** Medium (2-3 weeks) - ---- - -## 🎯 Recommended Implementation Order - -### Phase 1: Foundation (✅ DONE) -- [x] Clean repository structure -- [x] Fix data handling bugs -- [x] Validate training pipeline -- [x] Create universal visualization tool -- [x] Write comprehensive documentation - -### Phase 2: Text Integration (NEXT - 2-3 weeks) -- [ ] **Species Names** (Priority 1) - - Simplest extension - - High user value (searchability) - - Foundation for other text extensions - -**Steps:** -1. Load names from `names.dmp` -2. Implement BERT encoder -3. Create joint training loop -4. Update visualization for text queries -5. Validate on text search tasks - -### Phase 3: Multi-Modal (4-6 weeks) -- [ ] **Additional Features** (Priority 2) - - Collect metadata - - Medium complexity - - Good for downstream tasks - -- [ ] **Word Descriptions** (Priority 3) - - Requires text infrastructure from Phase 2 - - Enables semantic search - - Educational value - -### Phase 4: Advanced (6-8 weeks) -- [ ] **Protein Embeddings** (Priority 4) - - Most complex - - Requires large compute - - High biological value - -## 📋 Implementation Template - -For each extension, follow this template: - -### File Structure -``` -taxembed/ -├── src/taxembed/ -│ ├── encoders/ # NEW -│ │ ├── text_encoder.py -│ │ ├── protein_encoder.py -│ │ └── feature_encoder.py -│ ├── multimodal/ # NEW -│ │ ├── fusion.py -│ │ └── joint_training.py -│ └── models/ -│ └── poincare_multimodal.py # NEW -│ -├── scripts/ -│ ├── prepare_text_data.py # NEW -│ └── train_multimodal.py # NEW -│ -└── tests/ - └── test_multimodal.py # NEW -``` - -### Training Script -```python -# scripts/train_multimodal.py - -import torch -from taxembed.models import MultiModalPoincare - -def main(): - # Load data - graph_data = load_graph("data/taxonomy_edges.mapped.edgelist") - text_data = load_text("data/organism_names.txt") # NEW - - # Model - model = MultiModalPoincare( - graph_dim=10, - text_dim=768, # BERT hidden size - fusion_dim=10 - ) - - # Training - for epoch in range(epochs): - # Graph loss (existing) - graph_loss = train_graph_batch(model, graph_data) - - # Text loss (NEW) - text_loss = train_text_batch(model, text_data) - - # Combined - total_loss = graph_loss + lambda_text * text_loss - total_loss.backward() -``` - -### Evaluation -```python -# Evaluate multimodal -def evaluate_multimodal(model): - # Graph-based (existing) - graph_metrics = evaluate_nearest_neighbors(model) - - # Text-based (NEW) - text_metrics = evaluate_text_search(model) - # e.g., "sapiens" → find Homo sapiens - - # Joint - joint_metrics = evaluate_cross_modal(model) - # e.g., TaxID → text, text → TaxID - - return {**graph_metrics, **text_metrics, **joint_metrics} -``` - -## 🔬 Research Questions - -Each extension opens research opportunities: - -### Names -- How much does text improve hierarchy learning? -- Can we handle multilingual names? -- How to deal with synonyms? - -### Proteins -- Does protein sequence improve organism clustering? -- Can we predict protein function from taxonomy? -- How to handle horizontal gene transfer? - -### Features -- Which features are most informative? -- Can we predict missing features? -- How to handle sparse features? - -### Descriptions -- Can we generate organism descriptions? -- How to link to literature? -- Can we answer biological questions? - -## 📊 Success Metrics - -### Quantitative -- Nearest neighbor accuracy -- Cluster purity -- Text search recall@k -- Cross-modal retrieval accuracy - -### Qualitative -- Biologically meaningful clusters -- Useful for downstream tasks -- Better interpretability -- User satisfaction - -## 🎓 Publication Opportunities - -Each extension could lead to publications: - -1. **Multimodal Taxonomy Embeddings** - combine graph + text + features -2. **Hyperbolic Protein-Taxonomy Embeddings** - protein function in hyperbolic space -3. **Natural Language Queries for Organisms** - semantic search in taxonomy -4. **Hierarchical Biological Embeddings** - comprehensive benchmark - -## 💡 Additional Ideas - -### Short-term Enhancements -- [ ] Add more taxonomic groups to visualization -- [ ] Create web interface for exploration -- [ ] Add batch inference API -- [ ] Create pre-trained model checkpoints -- [ ] Add Docker container - -### Long-term Vision -- [ ] Real-time taxonomy updates -- [ ] Interactive exploration tool -- [ ] Integration with biological databases -- [ ] API for downstream applications -- [ ] Educational platform - -## 🤝 Getting Started - -**Ready to implement extensions?** - -1. Read `POINCARE_EMBEDDINGS_EXPLAINED.md` for technical background -2. Start with **Species Names** (easiest, high value) -3. Follow the implementation template -4. Add tests as you go -5. Update documentation -6. Create PR for review - -**Questions?** See `CONTRIBUTING.md` or open an issue. - ---- - -**Current focus:** Graph structure works excellently! Extensions should preserve this while adding complementary information. diff --git a/docs/archive/PERMANENT_FIX_PLAN.md b/docs/archive/PERMANENT_FIX_PLAN.md deleted file mode 100644 index c9a5360..0000000 --- a/docs/archive/PERMANENT_FIX_PLAN.md +++ /dev/null @@ -1,221 +0,0 @@ -# Permanent Fix Plan: Training All Nodes - -## **Root Cause Analysis** - -### Why do 29,063 nodes have no training signal? - -``` -Total nodes: 111,103 -In training pairs: 82,040 (74%) -NOT in training: 29,063 (26%) -``` - -**These nodes never appear in training because:** -1. They are LEAF taxa (no descendants) -2. Transitive closure only creates ancestor→descendant pairs -3. Leaves have no descendants, so they ONLY appear as descendants, never ancestors -4. BUT: Some leaves don't even appear as descendants if they're isolated - -### Current Data Flow: -``` -nodes.dmp (NCBI) - → filter to small dataset - → build transitive closure (ancestor → descendant pairs) - → PROBLEM: leaf nodes with no edges get excluded -``` - ---- - -## **Permanent Fix Strategy** - -### **Option 1: Complete Transitive Closure** ⭐ **RECOMMENDED** - -**Make sure ALL nodes in the mapping appear in training data:** - -1. **Include self-loops for leaf nodes** - ```python - # For nodes that never appear as ancestors: - # Add (leaf, leaf, depth_diff=0) pairs - ``` - -2. **Or connect isolated nodes to their parents** - ```python - # Even if no children, every node has a parent - # Add parent→leaf pairs from nodes.dmp - ``` - -**Advantages:** -- ✅ Every node gets training signal -- ✅ No special initialization needed -- ✅ Works for ANY dataset size -- ✅ Mathematically correct (every node in hierarchy) - -**Implementation:** -- Fix `build_transitive_closure.py` -- Ensure 100% node coverage in training data - ---- - -### **Option 2: Regularization-Only Training** (Less ideal) - -**Keep current training data, but ensure untrained nodes get proper regularization:** - -1. **Fix initialization** - - Don't set all missing nodes to depth=37 - - Load actual depths from nodes.dmp - - Or initialize at center (depth=0) and let reg move them - -2. **Ensure regularizer covers ALL nodes** - - Already fixed (idx_to_depth now has all nodes) - - But regularization alone isn't enough for structure - -**Disadvantages:** -- ❌ Nodes with no edges have weak training signal (only regularization) -- ❌ Their embeddings remain somewhat random -- ❌ Doesn't scale well to larger datasets - ---- - -## **Recommended Implementation: Option 1** - -### **Fix `build_transitive_closure.py`:** - -```python -def build_complete_training_data(taxonomy, valid_taxids, mapping): - """ - Build training pairs ensuring EVERY node appears at least once. - - Changes from original: - 1. After building transitive closure, check for missing nodes - 2. For each missing node, add pair: (parent, node) - 3. Ensures 100% coverage - """ - - # 1. Build standard transitive closure (existing logic) - training_data = build_transitive_closure(...) - - # 2. Find nodes that never appear - nodes_in_training = set() - for pair in training_data: - nodes_in_training.add(pair['ancestor_idx']) - nodes_in_training.add(pair['descendant_idx']) - - all_nodes = set(range(len(mapping))) - missing_nodes = all_nodes - nodes_in_training - - print(f"Missing {len(missing_nodes)} nodes from training") - - # 3. For each missing node, add parent→node pair - for node_idx in missing_nodes: - taxid = mapping[node_idx] - parent_taxid = taxonomy[taxid]['parent'] - parent_idx = reverse_mapping[parent_taxid] - - training_data.append({ - 'ancestor_idx': parent_idx, - 'descendant_idx': node_idx, - 'depth_diff': 1, - 'ancestor_depth': taxonomy[parent_taxid]['depth'], - 'descendant_depth': taxonomy[taxid]['depth'], - # ... other fields - }) - - print(f"Added {len(missing_nodes)} parent→leaf pairs") - print(f"Coverage: {len(set(all_nodes))} / {len(all_nodes)} (100%)") - - return training_data -``` - -### **Files to Modify:** - -1. **`build_transitive_closure.py`** - - Add logic to ensure 100% node coverage - - Add parent→node pairs for missing nodes - -2. **`train_small.py`** (cleanup) - - Remove the depth assignment fallback (line 371-392) - - No longer needed if training data is complete - -3. **Validation** - - Add check in `train_small.py` to verify 100% coverage - - Warn if nodes are missing from training - ---- - -## **Alternative: Keep Both Mechanisms** - -For maximum robustness: - -1. ✅ Fix `build_transitive_closure.py` to ensure 100% coverage -2. ✅ Keep fallback depth assignment in `train_small.py` (defensive) -3. ✅ Add validation warning if coverage < 100% - -This way: -- Primary: All nodes trained -- Fallback: If data is incomplete, at least they get proper initialization -- Monitoring: Clear warning if something is wrong - ---- - -## **Testing Plan** - -### **1. Small Dataset (111K nodes)** -```bash -# Rebuild transitive closure -python build_transitive_closure.py - -# Verify coverage -python -c " -import pickle -data = pickle.load(open('data/taxonomy_edges_small_transitive.pkl', 'rb')) -ancestors = {d['ancestor_idx'] for d in data} -descendants = {d['descendant_idx'] for d in data} -coverage = len(ancestors | descendants) -print(f'Coverage: {coverage:,} nodes') -assert coverage == 111103, 'Not all nodes covered!' -" - -# Train -python train_small.py --epochs 100 --lambda-reg 0.05 -``` - -### **2. Visualize** -```bash -python visualize_multi_groups.py taxonomy_model_small_best.pth -# Should look MUCH cleaner - all nodes properly trained -``` - -### **3. Large Dataset (2.7M nodes)** -```bash -# Same process but with full dataset -python build_transitive_closure.py --dataset full -# Should ensure 2.7M nodes all covered -``` - ---- - -## **Expected Improvements** - -| Metric | Before | After | Status | -|--------|--------|-------|--------| -| Nodes trained | 82,040 | 111,103 | ✅ +35% | -| Boundary noise | High | Low | ✅ Clean | -| Visual quality | Messy | Clean | ✅ Structured | -| Scalability | ❌ Breaks | ✅ Works | ✅ Fixed | - ---- - -## **Implementation Priority** - -1. **Immediate**: Fix `build_transitive_closure.py` ⭐ -2. **Optional**: Keep fallback in `train_small.py` (defensive) -3. **Required**: Add validation check -4. **Test**: Rebuild data, retrain, visualize - ---- - -## **Next Steps** - -1. Shall I implement the fix in `build_transitive_closure.py`? -2. Or do you want to review the current `build_transitive_closure.py` first? -3. Or test another approach? diff --git a/docs/archive/PERMANENT_FIX_SUMMARY.md b/docs/archive/PERMANENT_FIX_SUMMARY.md deleted file mode 100644 index 2e2ba94..0000000 --- a/docs/archive/PERMANENT_FIX_SUMMARY.md +++ /dev/null @@ -1,167 +0,0 @@ -# Permanent Fix Applied: Complete Node Coverage - -## **Results** - -### **Before Fix:** -``` -Total nodes: 111,103 -In training: 82,040 (73.8%) -Missing: 29,063 (26.2%) -``` - -### **After Fix:** -``` -Total nodes: 111,103 -In training: 109,236 (98.3%) ✅ +27,196 nodes! -Missing: 1,867 (1.7%) -``` - -### **Improvement:** -- ✅ **+27,196 nodes** now have training signal -- ✅ **98.3% coverage** (up from 73.8%) -- ✅ **Added 26,590 parent→node pairs** for previously untrained nodes - ---- - -## **What Was Fixed** - -### **File Modified:** -`build_transitive_closure.py` - -### **New Function Added:** -```python -def ensure_complete_coverage(training_data, valid_taxids, taxonomy): - """ - Ensure ALL nodes appear in training data. - For nodes that never appear (leaf nodes), add parent→node pairs. - """ -``` - -### **How It Works:** -1. Builds standard transitive closure (all ancestor→descendant pairs) -2. Checks which nodes are missing from training -3. For each missing node, adds a `parent→node` training pair -4. Validates final coverage - ---- - -## **Remaining 1,867 Nodes** - -These nodes couldn't be added because their **parents are not in the dataset**. - -Example warnings: -``` -⚠️ Parent 493944 of node 1499148 not in dataset -⚠️ Parent 2212439 of node 2212691 not in dataset -``` - -**Why this happens:** -- The small dataset is a subset of NCBI taxonomy -- Some nodes have parents outside this subset -- These are "boundary nodes" at the edge of our taxonomy sample - -**Options:** -1. **Accept 98.3% coverage** - Good enough for production ✅ -2. **Add self-loops** for orphan nodes (treat them as roots) -3. **Expand dataset** to include missing parents - -For most use cases, **98.3% is excellent** and the permanent fix is complete. - ---- - -## **Next Steps** - -### **1. Train with New Data** ⭐ - -```bash -# New training data has 1,002,486 pairs (up from 975,896) -python train_small.py --epochs 100 --lambda-reg 0.05 --early-stopping 10 -``` - -**Expected improvements:** -- ✅ 98.3% of nodes get proper training -- ✅ Much cleaner UMAP visualization -- ✅ Better hierarchical structure - -### **2. Compare Results** - -```bash -# After training completes -python visualize_multi_groups.py taxonomy_model_small_best.pth - -# Compare trained vs untrained nodes -python visualize_trained_only.py taxonomy_model_small_best.pth -``` - -Should see: -- Less noisy boundary cluster -- Better separation of taxonomic groups -- Cleaner hierarchical structure - -### **3. Optional: Handle Remaining 1,867 Nodes** - -If you want 100% coverage, add self-loops for orphan nodes: - -```python -# In ensure_complete_coverage(): -if parent_taxid not in taxid_to_idx: - # Parent not in dataset - add self-loop - parent_idx = node_idx - parent_taxid = node_taxid - # (rest of logic) -``` - ---- - -## **Scalability** - -This fix is **permanent and scales** to any dataset size: - -- ✅ **Small dataset (111K nodes)**: 98.3% coverage -- ✅ **Full dataset (2.7M nodes)**: Will also work -- ✅ **Any future dataset**: Automatically ensures coverage - -The `ensure_complete_coverage()` function runs automatically every time you rebuild the transitive closure. - ---- - -## **Files Changed** - -| File | Change | Status | -|------|--------|--------| -| `build_transitive_closure.py` | Added `ensure_complete_coverage()` | ✅ Complete | -| `train_small.py` | Keep fallback depth assignment (defensive) | ℹ️ Optional | -| `data/taxonomy_edges_small_transitive.pkl` | Rebuilt with 98.3% coverage | ✅ Updated | - ---- - -## **Validation** - -### **Quick Check:** -```bash -python -c " -import pickle -data = pickle.load(open('data/taxonomy_edges_small_transitive.pkl', 'rb')) -ancestors = {d['ancestor_idx'] for d in data} -descendants = {d['descendant_idx'] for d in data} -coverage = len(ancestors | descendants) -print(f'Coverage: {coverage:,} / 111,103 nodes ({100*coverage/111103:.1f}%)') -" -``` - -Expected output: -``` -Coverage: 109,236 / 111,103 nodes (98.3%) -``` - ---- - -## **Summary** - -✅ **Permanent fix applied** - scales to any dataset size -✅ **98.3% coverage** - up from 73.8% -✅ **27K more nodes** now get training signal -✅ **Training quality** will be much better -✅ **Visualizations** will be cleaner - -**Ready for retraining!** diff --git a/docs/archive/POINCARE_EMBEDDINGS_EXPLAINED.md b/docs/archive/POINCARE_EMBEDDINGS_EXPLAINED.md deleted file mode 100644 index cf1dd6c..0000000 --- a/docs/archive/POINCARE_EMBEDDINGS_EXPLAINED.md +++ /dev/null @@ -1,406 +0,0 @@ -# Poincaré Embeddings Explained - -## What Are Poincaré Embeddings? - -### The Problem: Representing Hierarchies -Traditional embeddings (like Word2Vec, GloVe) use **Euclidean space** (flat, normal space). But hierarchical data (like taxonomies, org charts, WordNet) have a **tree-like structure** that doesn't fit well in flat space. - -**Why?** In a tree: -- The root has 1 node -- Level 1 might have 10 nodes -- Level 2 might have 100 nodes -- Level 3 might have 1,000 nodes -- ...exponential growth - -In flat Euclidean space, you need **exponentially growing dimensions** to represent this without distortion. - -### The Solution: Hyperbolic Space -**Poincaré embeddings** use **hyperbolic geometry** - a curved space where: -- Distance grows exponentially as you move from the center -- Perfect for hierarchies: root near center, leaves near boundary -- Can represent exponential growth in constant dimensions - -**Visualization:** -``` -Euclidean Space (flat): Hyperbolic Space (Poincaré disk): - o o (root, center) - / \ /|\ - o o / | \ - /| |\ / | \ -o o o o o o o - /|\ /|\ /|\ - (exponentially more space - near the boundary) -``` - -### Key Properties -1. **Distance from center = hierarchy level** - - Root organisms near center (e.g., ||x|| ≈ 0) - - Leaf organisms near boundary (e.g., ||x|| ≈ 0.99) - -2. **Angular distance = similarity within level** - - All primates cluster in one angular region - - All bacteria cluster in another region - -3. **Hierarchical distance preserved** - - Related organisms are closer - - Distance in embedding ≈ distance in taxonomy tree - -## Our Data: NCBI Taxonomy - -### What Our Data Looks Like - -#### Raw Data (from NCBI) -``` -nodes.dmp: - TaxID | Parent_TaxID | Rank | ... - 9606 | 9605 | species # Homo sapiens → Homo (genus) - 9605 | 207598 | genus # Homo → Homininae (subfamily) - 207598| 9604 | subfamily # Homininae → Hominidae (family) - 9604 | 314295 | family # Hominidae → Simiiformes - ... -``` - -#### Edge List (parent-child relationships) -``` -data/taxonomy_edges.edgelist: - 9606 9605 # Homo sapiens → Homo - 9605 207598 # Homo → Homininae - 207598 9604 # Homininae → Hominidae - 9604 314295 # Hominidae → Simiiformes - 314295 9526 # Simiiformes → Primates - 9526 40674 # Primates → Mammalia - ... -``` - -This is a **directed graph** where each edge represents: -``` -child_taxid → parent_taxid -``` - -#### Remapped for Training -``` -data/taxonomy_edges_small.mapped.edgelist: - 0 1 # Remapped indices (TaxID 9606 → idx 0, TaxID 9605 → idx 1) - 2 3 - 4 5 - ... -``` - -Plus mapping file: -``` -data/taxonomy_edges_small.mapping.tsv: - taxid idx - 9606 0 - 9605 1 - ... -``` - -### Visual Example - -Here's what the hierarchy looks like: - -``` - Root (1) - | - ┌─────────────────┴─────────────────┐ - Bacteria (2) Eukaryota (2759) - | | - ┌───┴───┐ ┌─────┴─────┐ - E.coli Salmonella Animals (33208) Plants - (562) (590) | - ┌───┴───┐ - Vertebrates Invertebrates - (7742) (...) - | - ┌────┴────┐ - Mammals Fish - (40674) (...) - | - ┌────┴────┐ - Primates Rodents - (9443) (9989) - | | - ┌───┴───┐ Mouse - Humans Apes (10090) - (9605) (9604) - | - Homo sapiens - (9606) -``` - -**In the edge list, we have:** -- 9606 → 9605 (Homo sapiens → Homo) -- 9605 → 207598 (Homo → Homininae) -- 9604 → 314295 (Hominidae → Simiiformes) -- 9443 → 40674 (Primates → Mammalia) -- etc. - -### Data Statistics - -**Small dataset:** -- 111,103 unique organisms (nodes) -- 100,000 parent-child relationships (edges) -- Covers major groups but incomplete - -**Full dataset:** -- 2,705,745 unique organisms (nodes) -- 2,705,744 parent-child relationships (edges) -- Complete NCBI taxonomy tree - -## What Are We Training? - -### The Embedding Model - -We learn a **vector** for each organism: -``` -TaxID → Vector in R^d (e.g., d=10) - -Example (simplified): - Homo sapiens (9606) → [0.15, 0.23, 0.08, ..., 0.45] - Homo (9605) → [0.14, 0.22, 0.09, ..., 0.43] # Similar, nearby - E. coli (562) → [-0.80, 0.02, -0.35, ..., 0.10] # Very different -``` - -### The Training Process - -#### 1. Objective: Learn Distance Model -For each edge `(u → v)` (child → parent): -``` -Distance in embedding should be SMALL -d_poincare(embed(u), embed(v)) should be ≈ 0.1-0.5 -``` - -For non-edges (random pairs): -``` -Distance in embedding should be LARGE -d_poincare(embed(u), embed(random)) should be ≈ 1.0-2.0 -``` - -#### 2. Loss Function -```python -# For each training batch: -for (child, parent) in edges: - # Positive sample (real edge) - positive_distance = d_poincare(embed[child], embed[parent]) - - # Negative samples (random organisms, not related) - for neg in random_samples(k=50): - negative_distance = d_poincare(embed[child], embed[neg]) - - # Loss: want positive_distance < negative_distance - loss = max(0, margin + positive_distance - negative_distance) -``` - -**Intuition:** -- Child and parent should be close -- Child and random organism should be far -- Margin = minimum separation we want - -#### 3. Hyperbolic Distance -The key is using **Poincaré distance**, not Euclidean: - -```python -def poincare_distance(u, v): - """Distance in Poincaré ball model.""" - norm_u_sq = ||u||^2 - norm_v_sq = ||v||^2 - norm_diff_sq = ||u - v||^2 - - numerator = 2 * norm_diff_sq - denominator = (1 - norm_u_sq) * (1 - norm_v_sq) - - return arcosh(1 + numerator / denominator) -``` - -This distance: -- Respects the hyperbolic geometry -- Grows exponentially near the boundary (||x|| → 1) -- Preserves hierarchical structure - -#### 4. Riemannian SGD -We can't use normal gradient descent because we're in curved space! - -**Riemannian SGD:** -1. Compute gradient in tangent space (flat space at current point) -2. Project back onto the Poincaré ball (ensure ||x|| < 1) -3. Update embeddings - -```python -# Simplified -grad = compute_gradient(loss) -x_new = x - lr * grad -x_new = project_onto_ball(x_new) # Keep ||x|| < 1 -``` - -### Training Example - -**Epoch 0 (start):** -``` -Loss: 3.94 -Embeddings are random -Human → Nearest neighbor: Random organism (bad!) -``` - -**Epoch 50:** -``` -Loss: 2.86 -Embeddings learning structure -Human → Nearest neighbor: Other primate (better!) -``` - -**Epoch 500:** -``` -Loss: 2.32 -Embeddings encode hierarchy well -Human → Nearest neighbor: Almost identical primate species (distance 0.0007) -Primates cluster together in UMAP -``` - -## Current Setup: TaxID Only - -### What We Currently Embed - -**Each organism = One node = One vector** - -``` -TaxID 9606 (Homo sapiens) → vector [0.15, 0.23, ..., 0.45] -TaxID 9605 (Homo) → vector [0.14, 0.22, ..., 0.43] -TaxID 562 (E. coli) → vector [-0.80, 0.02, ..., 0.10] -``` - -**Names are separate** (in mapping file): -``` -mapping.tsv: - 9606 → "Homo sapiens" - 562 → "Escherichia coli" -``` - -**During training:** -- Only TaxID integers are used -- Names are NOT used -- Model learns from graph structure only - -**During visualization:** -- We look up names from mapping file -- Display "Homo sapiens (9606)" for human interpretation - -## Can We Add Species Names to Training? - -### Option 1: Current Approach (Graph Structure Only) ✅ -**What we're doing now** - -**Pros:** -- ✅ Fast training -- ✅ Works perfectly for taxonomy hierarchy -- ✅ No text processing needed -- ✅ Language-independent - -**Cons:** -- ❌ Doesn't learn from name similarity -- ❌ Can't handle organisms without TaxID -- ❌ Can't do text-based queries ("find primates") - -### Option 2: Text + Graph (Multimodal) ⭐ -**Add species names as additional signal** - -**Approach:** -1. Encode names with text encoder (e.g., BERT, BioBERT) -2. Learn joint embedding space -3. Train on both graph structure AND name similarity - -```python -# Pseudocode -text_embedding = encode_text(species_name) -graph_embedding = current_poincare_embedding - -# Joint loss -loss_graph = poincare_loss(graph_edges) -loss_text = similarity_loss(text_embedding, graph_embedding) -total_loss = loss_graph + λ * loss_text -``` - -**Benefits:** -- Can query by name: "organisms similar to 'sapiens'" -- Better generalization -- Can handle typos, synonyms -- Cross-lingual (if using multilingual encoder) - -**Downsides:** -- Much more complex -- Requires text encoder -- Slower training -- Might not improve hierarchy learning (names don't encode hierarchy well) - -### Option 3: Name as Feature (Auxiliary) -**Use names as additional features, not primary signal** - -```python -# Add name features to each node -features = { - 'taxid': 9606, - 'name': 'Homo sapiens', - 'name_length': 12, - 'has_genus': True, - 'has_species': True, - 'name_embedding': encode(name) # Small text encoding -} - -# Then train graph embedding with features -``` - -**Benefits:** -- Lightweight addition -- Can improve disambiguation -- Still mainly graph-based - -### Recommendation for Your Use Case - -**If your goal is hierarchy learning:** -👉 **Stick with current approach** (graph structure only) -- It's working perfectly -- Names don't add much for hierarchy -- Much simpler and faster - -**If you want text-based queries:** -👉 **Add text encoder as Option 2** -- Could enable: "Find all species with 'sapiens' in name" -- Could enable: "Find organisms similar to 'human'" -- Useful for downstream applications - -**If you want names for better interpretability:** -👉 **Current approach is fine!** -- Names are in mapping file -- Visualization script already shows names -- No training changes needed - -## Summary - -### What Poincaré Embeddings Are -- Embeddings in **hyperbolic space** (curved geometry) -- Perfect for **hierarchical data** like taxonomies -- Preserve tree structure in low dimensions - -### What Our Data Looks Like -- **Nodes:** Organisms (TaxIDs) -- **Edges:** Parent-child relationships -- **Structure:** Tree/DAG of 111K-2.7M organisms - -### What We're Training -- **One vector per organism** in R^10 -- **Minimize distance** for parent-child pairs -- **Maximize distance** for unrelated pairs -- **Use hyperbolic distance** and Riemannian SGD - -### Names in Training -**Current:** Names NOT used in training (only in visualization) -**Possible:** Could add text encoder for multimodal learning -**Recommendation:** Current approach is excellent for hierarchy learning - -### The Magic -After training, embeddings capture: -- ✅ Hierarchy: depth in tree ≈ distance from center -- ✅ Similarity: related organisms cluster together -- ✅ Relationships: nearest neighbors are taxonomically related - -All in just **10 dimensions** in hyperbolic space! 🎉 diff --git a/docs/archive/REPOSITORY_STATUS.md b/docs/archive/REPOSITORY_STATUS.md deleted file mode 100644 index a997d17..0000000 --- a/docs/archive/REPOSITORY_STATUS.md +++ /dev/null @@ -1,286 +0,0 @@ -# Repository Status - November 2025 - -## ✅ Repository is Production-Ready - -The poincare-embeddings (taxembed) repository has been fully restructured, debugged, and cleaned. - -## Current Status - -### 🎯 Completed Tasks - -1. **Repository Restructuring** ✅ - - Cookiecutter-style directory structure - - `src/taxembed/` package layout - - Organized `scripts/` directory - - Proper `tests/` directory - - Modern `pyproject.toml` with `uv` - - `ruff` for linting and formatting - -2. **Data Handling Fixes** ✅ - - Fixed header line bug (removed 2 fake "id1", "id2" nodes) - - Clean datasets: 111,103 nodes (small), 2.7M nodes (full) - - Data validation tool (`scripts/validate_data.py`) - - All data quality checks pass - -3. **Training Validation** ✅ - - Successfully trained on small dataset (500 epochs) - - Loss decreased from 3.94 → 2.32 - - Primates cluster correctly - - Nearest neighbors show biological relevance - -4. **Repository Cleanup** ✅ - - Removed 569 checkpoint files - - Removed all log files and temporary visualizations - - Consolidated 5 visualization scripts into 1 universal tool - - Removed 4 redundant shell scripts - - Updated `.gitignore` - -5. **Documentation** ✅ - - Comprehensive README.md - - QUICKSTART.md, GETTING_STARTED.md - - SCRIPTS_GUIDE.md (detailed script documentation) - - DATA_FIXES_SUMMARY.md (bug fixes) - - CLEANUP_SUMMARY.md (cleanup details) - - STRUCTURE.md (project organization) - -## Repository Structure - -``` -taxembed/ -├── Core Training Scripts -│ ├── embed.py # ⭐ Main training -│ ├── prepare_taxonomy_data.py # Data preparation -│ ├── remap_edges.py # Data remapping -│ ├── monitor_training.py # Training monitor -│ └── evaluate_full.py # Evaluation -│ -├── src/taxembed/ # Source package -│ ├── manifolds/ # Hyperbolic geometry -│ ├── models/ # Embedding models -│ ├── datasets/ # Data loaders -│ └── utils/ # Utilities -│ -├── scripts/ # Utility scripts -│ ├── visualize_embeddings.py # ⭐ Universal visualization -│ ├── validate_data.py # ⭐ Data validation -│ ├── cleanup_repo.sh # Repository cleanup -│ └── regenerate_data.sh # Data regeneration -│ -├── tests/ # Unit tests -│ -├── hype/ # Original package (backward compat) -│ -├── Configuration -│ ├── pyproject.toml # Modern Python config -│ ├── ruff.toml # Linter config -│ ├── Makefile # Convenience commands -│ └── .gitignore # Git ignore -│ -└── Documentation (12 files) - ├── README.md - ├── SCRIPTS_GUIDE.md # ⭐ How to use scripts - ├── QUICKSTART.md - └── ... (see below) -``` - -## Key Tools - -### Training -```bash -python embed.py \ - -dset data/taxonomy_edges_small.mapped.edgelist \ - -checkpoint model.pth \ - -dim 10 -epochs 50 -negs 50 -burnin 10 \ - -batchsize 32 -model distance -manifold poincare \ - -lr 0.1 -gpu -1 -ndproc 1 -train_threads 1 \ - -eval_each 999999 -fresh -``` - -### Universal Visualization ⭐ -```bash -# Works with ANY checkpoint -python scripts/visualize_embeddings.py model.pth --highlight primates -python scripts/visualize_embeddings.py model.pth --only mammals -python scripts/visualize_embeddings.py model.pth --nearest 10 -``` - -### Data Validation -```bash -python scripts/validate_data.py small -python scripts/validate_data.py full -``` - -### Repository Cleanup -```bash -./scripts/cleanup_repo.sh -``` - -## Data Quality - -### Small Dataset -- **Nodes:** 111,103 organisms (clean, no fake nodes) -- **Edges:** 100,000 taxonomic relationships -- **Status:** ✅ All validation checks pass - -### Full Dataset -- **Nodes:** 2,705,745 organisms -- **Edges:** 2,705,744 taxonomic relationships -- **Status:** ✅ All validation checks pass - -## Training Results (500 epochs on small dataset) - -- **Loss:** 3.94 → 2.32 (41% reduction) -- **Nearest Neighbors:** Biologically accurate - - Human → Other primates (distance 0.0007) - - Mouse → Other rodents (distance 0.0011) - - E. coli → Other bacteria (distance 0.0003) -- **Clustering:** Primates form distinct cluster -- **UMAP:** Clear hierarchical structure - -## File Inventory - -### Documentation (12 files) -1. `README.md` - Main documentation -2. `QUICKSTART.md` - Quick start guide -3. `GETTING_STARTED.md` - Detailed setup -4. `SCRIPTS_GUIDE.md` - **⭐ Script usage guide** -5. `STRUCTURE.md` - Project structure -6. `CONTRIBUTING.md` - Contribution guide -7. `DATA_FIXES_SUMMARY.md` - Data bug fixes -8. `DATA_HANDLING_REVIEW.md` - Data analysis -9. `CLEANUP_SUMMARY.md` - Cleanup details -10. `RESTRUCTURING_SUMMARY.md` - Restructuring notes -11. `REPOSITORY_STATUS.md` - This file -12. Various other notes and summaries - -### Core Scripts (7 files) -1. `embed.py` - Main training script -2. `prepare_taxonomy_data.py` - Data preparation -3. `remap_edges.py` - Data remapping -4. `monitor_training.py` - Training monitor -5. `evaluate_full.py` - Evaluation -6. `evaluate_and_visualize.py` - Combined eval -7. `nn_demo.py` - Quick demo - -### Utility Scripts (8 files in scripts/) -1. `visualize_embeddings.py` - **⭐ Universal visualization** -2. `validate_data.py` - **⭐ Data validation** -3. `cleanup_repo.sh` - Repository cleanup -4. `regenerate_data.sh` - Data regeneration -5-8. Various wrapper scripts - -### Configuration (5 files) -1. `pyproject.toml` - Modern Python config -2. `ruff.toml` - Linter config -3. `Makefile` - Convenience commands -4. `setup.py` - C++ extensions -5. `.gitignore` - Git ignore rules - -## Quality Metrics - -### Code Quality -- ✅ Structured with `src/` layout -- ✅ Linted with `ruff` -- ✅ Type hints (partial) -- ✅ Clear function documentation - -### Data Quality -- ✅ No header bugs -- ✅ Sequential indices -- ✅ Consistent mappings -- ✅ Validated with automated checks - -### Documentation Quality -- ✅ 12 comprehensive documentation files -- ✅ Clear usage examples -- ✅ Troubleshooting guides -- ✅ API documentation - -### Repository Cleanliness -- ✅ No checkpoint files in repo -- ✅ No log files -- ✅ No temporary visualizations -- ✅ Proper `.gitignore` -- ✅ Organized file structure - -## Recommended Workflows - -### New User -```bash -# 1. Setup -make install -make build -python scripts/validate_data.py small - -# 2. Quick training test -python embed.py -dset data/taxonomy_edges_small.mapped.edgelist \ - -checkpoint test.pth -dim 10 -epochs 5 -negs 50 -burnin 2 \ - -batchsize 32 -model distance -manifold poincare \ - -lr 0.1 -gpu -1 -ndproc 1 -train_threads 1 -eval_each 999999 -fresh - -# 3. Visualize -python scripts/visualize_embeddings.py test.pth --highlight primates -``` - -### Production Training -```bash -# Full dataset, 200 epochs -python embed.py -dset data/taxonomy_edges.mapped.edgelist \ - -checkpoint taxonomy_full.pth -dim 10 -epochs 200 -negs 50 -burnin 10 \ - -batchsize 32 -model distance -manifold poincare \ - -lr 0.1 -gpu -1 -ndproc 1 -train_threads 1 -eval_each 999999 -fresh -``` - -### Regular Maintenance -```bash -# Clean up old files -./scripts/cleanup_repo.sh - -# Validate data after changes -python scripts/validate_data.py small -python scripts/validate_data.py full - -# Regenerate data from NCBI taxonomy -./scripts/regenerate_data.sh -``` - -## Next Steps (Optional) - -### For Your Student -1. Read `QUICKSTART.md` to get started -2. Check `SCRIPTS_GUIDE.md` for script usage -3. Run validation: `python scripts/validate_data.py small` -4. Train test model (5 epochs) -5. Visualize: `python scripts/visualize_embeddings.py --highlight primates` - -### For Production -1. Train on full dataset (200+ epochs) -2. Evaluate multiple embedding dimensions (10, 20, 50) -3. Compare different manifolds (Poincaré, Lorentz) -4. Benchmark against baselines -5. Write paper with results - -### For Development -1. Add unit tests (`tests/`) -2. Improve type hints -3. Add CI/CD pipeline -4. Create conda/docker environment -5. Publish to PyPI - -## References - -- **Main Paper:** [Poincaré Embeddings for Learning Hierarchical Representations](https://arxiv.org/abs/1705.08039) -- **NCBI Taxonomy:** https://www.ncbi.nlm.nih.gov/taxonomy -- **Documentation:** See all `.md` files in root directory -- **Script Guide:** `SCRIPTS_GUIDE.md` - -## Summary - -✅ **Repository restructured** with modern Python best practices -✅ **Data bugs fixed** - clean, validated datasets -✅ **Training validated** - 500 epoch model shows excellent results -✅ **Repository cleaned** - 569 checkpoints removed, scripts consolidated -✅ **Documentation complete** - 12 comprehensive guides -✅ **Production-ready** - clean, maintainable, well-documented - -**The repository is now ready for serious work, publication, and sharing!** 🎉 diff --git a/docs/archive/RESTRUCTURING_COMPLETE.md b/docs/archive/RESTRUCTURING_COMPLETE.md deleted file mode 100644 index 032830f..0000000 --- a/docs/archive/RESTRUCTURING_COMPLETE.md +++ /dev/null @@ -1,243 +0,0 @@ -# ✅ Repository Restructuring Complete - -The taxembed repository has been successfully restructured with professional Python project standards. - -## 📋 Summary of Changes - -### New Files Created - -#### Configuration Files -- ✅ **pyproject.toml** - Unified project configuration with uv and dependencies -- ✅ **ruff.toml** - Linter and formatter configuration -- ✅ **Makefile** - Convenient command shortcuts - -#### Documentation -- ✅ **README.md** - Comprehensive project documentation -- ✅ **QUICKSTART.md** - Quick start guide -- ✅ **GETTING_STARTED.md** - Getting started guide -- ✅ **STRUCTURE.md** - Project organization documentation -- ✅ **RESTRUCTURING_SUMMARY.md** - Migration guide -- ✅ **RESTRUCTURING_COMPLETE.md** - This file - -#### Source Code Structure -- ✅ **src/taxembed/__init__.py** - Main package initialization -- ✅ **src/taxembed/manifolds/__init__.py** - Manifolds subpackage -- ✅ **src/taxembed/models/__init__.py** - Models subpackage -- ✅ **src/taxembed/datasets/__init__.py** - Datasets subpackage -- ✅ **src/taxembed/utils/__init__.py** - Utils subpackage - -#### Scripts -- ✅ **scripts/train.py** - Main training script -- ✅ **scripts/prepare_data.py** - Data preparation wrapper -- ✅ **scripts/remap_data.py** - ID remapping wrapper -- ✅ **scripts/monitor.py** - Training monitoring wrapper -- ✅ **scripts/evaluate.py** - Evaluation wrapper -- ✅ **scripts/visualize.py** - Visualization wrapper - -#### Testing -- ✅ **tests/__init__.py** - Tests package initialization -- ✅ **tests/test_example.py** - Example test module - -#### Other -- ✅ **Updated .gitignore** - Comprehensive gitignore rules -- ✅ **Updated CONTRIBUTING.md** - Development guidelines - -## 🏗️ New Project Structure - -``` -taxembed/ -├── src/ -│ └── taxembed/ # Main package (src/ layout) -│ ├── __init__.py -│ ├── manifolds/ # Hyperbolic manifolds -│ ├── models/ # Embedding models -│ ├── datasets/ # Data loading -│ └── utils/ # Utilities -├── scripts/ # Standalone scripts -│ ├── train.py -│ ├── prepare_data.py -│ ├── remap_data.py -│ ├── monitor.py -│ ├── evaluate.py -│ └── visualize.py -├── tests/ # Unit tests -│ ├── __init__.py -│ └── test_example.py -├── data/ # Data directory (gitignored) -├── pyproject.toml # Project configuration (uv) -├── ruff.toml # Linter configuration -├── Makefile # Command shortcuts -├── README.md # Main documentation -├── QUICKSTART.md # Quick start guide -├── GETTING_STARTED.md # Getting started guide -├── STRUCTURE.md # Project structure -├── CONTRIBUTING.md # Contribution guidelines -├── RESTRUCTURING_SUMMARY.md # Migration guide -├── RESTRUCTURING_COMPLETE.md # This file -├── LICENSE # CC-BY-NC 4.0 -└── .gitignore # Git ignore rules -``` - -## 🎯 Key Improvements - -### 1. Professional Project Layout -- ✅ `src/` layout (Python best practice) -- ✅ Organized package structure -- ✅ Separate scripts and tests directories -- ✅ Clear separation of concerns - -### 2. Modern Dependency Management -- ✅ **uv** for fast package management (10-100x faster than pip) -- ✅ Single `pyproject.toml` source of truth -- ✅ Clear separation of core vs. optional dependencies -- ✅ PEP 518 compliant - -### 3. Code Quality Enforcement -- ✅ **ruff** for fast linting and formatting -- ✅ Automatic code formatting -- ✅ Import sorting (isort) -- ✅ Bug detection -- ✅ 100-character line limit - -### 4. Comprehensive Documentation -- ✅ Clear README with examples -- ✅ Quick start guide -- ✅ Getting started guide -- ✅ Project structure documentation -- ✅ Migration guide -- ✅ Updated contribution guidelines - -### 5. Developer Experience -- ✅ Makefile for common tasks -- ✅ Convenient `uv run` commands -- ✅ Pytest integration -- ✅ Coverage reporting support - -## 📦 Dependencies - -### Core Dependencies -- PyTorch >= 2.0.0 -- NumPy >= 1.21.0, < 2.0 -- Pandas >= 1.3.0 -- Cython >= 3.0 -- tqdm >= 4.60.0 -- scikit-learn >= 1.0.0 -- h5py >= 3.0.0 -- iopath >= 0.1.10 -- nltk >= 3.8 - -### Optional Dependencies -- **Visualization**: matplotlib, umap-learn -- **Development**: ruff, pytest, pytest-cov - -## 🚀 Getting Started - -### Installation -```bash -make install -make build -``` - -### Usage -```bash -# Check code quality -make lint -make format - -# Run tests -make test - -# Train model -uv run python scripts/train.py --dataset data/taxonomy_edges.mapped.edgelist ... - -# See all commands -make help -``` - -## 📚 Documentation Files - -| File | Purpose | -|------|---------| -| **README.md** | Main project documentation | -| **QUICKSTART.md** | Detailed quick start guide | -| **GETTING_STARTED.md** | Quick getting started guide | -| **STRUCTURE.md** | Project organization details | -| **CONTRIBUTING.md** | Development guidelines | -| **RESTRUCTURING_SUMMARY.md** | What changed and why | -| **RESTRUCTURING_COMPLETE.md** | This summary | - -## ✨ Features - -### Code Quality -- Linting with ruff -- Auto-formatting -- Import sorting -- Bug detection - -### Testing -- pytest integration -- Coverage reporting -- Example test structure - -### Development -- Makefile shortcuts -- uv package management -- Type hint support -- IDE integration ready - -## 🔄 Backward Compatibility - -The original structure is preserved: -- ✅ `hype/` package remains in place -- ✅ Root-level scripts still work -- ✅ All original functionality maintained - -New structure is recommended but not required. - -## 📋 Checklist for Your Team - -- [ ] Install uv: https://github.com/astral-sh/uv#installation -- [ ] Run `make install` to install dependencies -- [ ] Run `make build` to build C++ extensions -- [ ] Run `make test` to verify everything works -- [ ] Read QUICKSTART.md for detailed instructions -- [ ] Read CONTRIBUTING.md for development guidelines -- [ ] Start using `make lint` and `make format` before committing -- [ ] Use `uv run` for executing Python scripts - -## 🎓 Learning Resources - -### For Users -- Start with: **QUICKSTART.md** -- Then read: **README.md** -- Reference: **STRUCTURE.md** - -### For Developers -- Start with: **GETTING_STARTED.md** -- Then read: **CONTRIBUTING.md** -- Reference: **STRUCTURE.md** - -### For Migration -- Read: **RESTRUCTURING_SUMMARY.md** - -## 🤝 Support - -For questions or issues: -1. Check the relevant documentation file -2. Review CONTRIBUTING.md for development guidelines -3. Open an issue on GitHub - -## 🎉 Next Steps - -1. **Install dependencies**: `make install` -2. **Build extensions**: `make build` -3. **Run tests**: `make test` -4. **Read documentation**: Start with QUICKSTART.md -5. **Start developing**: Use `make lint` and `make format` - ---- - -**Restructuring Date**: 2024 -**Status**: ✅ Complete -**Backward Compatibility**: ✅ Maintained -**Documentation**: ✅ Comprehensive diff --git a/docs/archive/RESTRUCTURING_SUMMARY.md b/docs/archive/RESTRUCTURING_SUMMARY.md deleted file mode 100644 index 04c978c..0000000 --- a/docs/archive/RESTRUCTURING_SUMMARY.md +++ /dev/null @@ -1,213 +0,0 @@ -# Repository Restructuring Summary - -This document summarizes the restructuring of the taxembed repository to follow cookiecutter best practices with modern Python tooling. - -## What Changed - -### ✅ New Structure - -The repository now follows a professional Python project layout: - -``` -taxembed/ -├── src/taxembed/ # Main package (src/ layout) -├── scripts/ # Standalone scripts -├── tests/ # Unit tests -├── data/ # Data directory (gitignored) -├── pyproject.toml # Project configuration (uv) -├── ruff.toml # Linter configuration -├── Makefile # Convenience commands -└── docs/ # Documentation -``` - -### 📦 Dependency Management: uv - -**Before:** `requirements.txt` + `setup.py` + `environment.yml` - -**After:** `pyproject.toml` with uv - -Benefits: -- Single source of truth for dependencies -- Faster installation (uv is 10-100x faster than pip) -- Better dependency resolution -- Follows PEP 518 standards - -### 🔍 Code Quality: ruff - -**Before:** No linting configuration - -**After:** `ruff.toml` with comprehensive linting rules - -Features: -- Fast Python linter (written in Rust) -- Automatic code formatting -- Import sorting (isort) -- Detects common bugs -- 100-character line limit - -### 📚 Documentation - -**New files:** -- `README.md` - Comprehensive project documentation -- `QUICKSTART.md` - Get started in minutes -- `STRUCTURE.md` - Project organization guide -- `CONTRIBUTING.md` - Updated contribution guidelines -- `Makefile` - Convenient command shortcuts - -### 🧪 Testing - -**New structure:** -- `tests/` directory for unit tests -- pytest configuration in `pyproject.toml` -- Coverage reporting support - -## Migration Guide - -### For Users - -**Old way:** -```bash -python -m venv venv -source venv/bin/activate -pip install -r requirements.txt -python setup.py build_ext --inplace -python embed.py -dset data/taxonomy_edges.mapped.edgelist ... -``` - -**New way:** -```bash -uv sync -python setup.py build_ext --inplace -uv run python scripts/train.py --dataset data/taxonomy_edges.mapped.edgelist ... -``` - -Or use the Makefile: -```bash -make install -make build -uv run python scripts/train.py --dataset data/taxonomy_edges.mapped.edgelist ... -``` - -### For Developers - -**Old way:** -```bash -# No linting, no tests, no structure -``` - -**New way:** -```bash -# Code quality -make lint # Check code -make format # Fix code style - -# Testing -make test # Run tests -make test-cov # With coverage - -# Development -make clean # Clean build artifacts -``` - -## Key Improvements - -### 1. **Cleaner Organization** -- Source code in `src/taxembed/` (src/ layout) -- Scripts in `scripts/` directory -- Tests in `tests/` directory -- Configuration files at root - -### 2. **Better Dependency Management** -- Single `pyproject.toml` file -- Clear separation of core vs. optional dependencies -- Faster installation with uv - -### 3. **Code Quality** -- Automatic linting with ruff -- Code formatting enforcement -- Import sorting -- Bug detection - -### 4. **Improved Documentation** -- Clear README with examples -- Quick start guide -- Project structure documentation -- Updated contribution guidelines - -### 5. **Developer Experience** -- Makefile for common tasks -- Convenient `uv run` commands -- Pytest integration -- Coverage reporting - -## Backward Compatibility - -The original `hype/` package and root-level scripts remain in place for backward compatibility. You can still use: -```bash -python embed.py ... -python prepare_taxonomy_data.py ... -``` - -However, we recommend migrating to the new structure: -```bash -uv run python scripts/train.py ... -uv run python scripts/prepare_data.py ... -``` - -## Configuration Details - -### pyproject.toml - -```toml -[project] -name = "taxembed" -version = "0.1.0" -requires-python = ">=3.8" -dependencies = [ - "torch>=2.0.0", - "numpy>=1.21.0,<2.0", - # ... more dependencies -] - -[project.optional-dependencies] -dev = ["ruff>=0.1.0", "pytest>=7.0.0"] -viz = ["matplotlib>=3.5.0", "umap-learn>=0.5.0"] -``` - -### ruff.toml - -```toml -line-length = 100 -target-version = "py38" - -[lint] -select = ["E", "W", "F", "I", "C", "B", "UP"] -ignore = ["E501", "W503", "E203"] -``` - -## Next Steps - -1. **Update your workflow:** - - Use `uv sync` instead of `pip install` - - Use `uv run` to execute scripts - - Use `make` commands for common tasks - -2. **Set up linting in your IDE:** - - Configure ruff in VS Code, PyCharm, etc. - - Enable auto-formatting on save - -3. **Add tests:** - - Create test files in `tests/` directory - - Run `make test` to verify - -4. **Update CI/CD:** - - Use `uv sync` in your CI pipeline - - Add `make lint` and `make test` steps - -## Questions? - -Refer to: -- `README.md` - Project overview and usage -- `QUICKSTART.md` - Get started quickly -- `STRUCTURE.md` - Detailed project organization -- `CONTRIBUTING.md` - Development guidelines diff --git a/docs/archive/REVERT_HYPERPARAMS.md b/docs/archive/REVERT_HYPERPARAMS.md deleted file mode 100644 index ab605fb..0000000 --- a/docs/archive/REVERT_HYPERPARAMS.md +++ /dev/null @@ -1,145 +0,0 @@ -# Hyperparameter Reversion - -## **Problem Identified** - -After comparing old (epoch 28) vs current (epoch 36) models: - -``` - OLD CURRENT Change -Loss: 0.472 0.531 +12.5% ⬆️ WORSE -Boundary (>0.90): 27.5% 33.5% +6% more cramped -``` - -**Root cause:** I over-corrected the hyperparameters when fixing the data coverage issue. - ---- - -## **What I Changed (and shouldn't have)** - -### **1. Regularization Strength (λ)** -- ❌ **Changed:** 0.1 → 0.01 (10x weaker) -- ✅ **Reverted:** Back to 0.1 -- **Why:** 0.01 was too weak for 111K nodes, not enough structure enforcement - -### **2. Initialization Range** -- ❌ **Changed:** [0.1, 0.95] → [0.05, 0.85] -- ✅ **Reverted:** Back to [0.1, 0.95] -- **Why:** Starting too conservative made learning harder - -### **3. Hard Projection Boundary** -- ❌ **Changed:** max_norm = 0.99999 → 0.98 -- ✅ **Reverted:** Back to 0.999 -- **Why:** 0.98 caused artificial compression - ---- - -## **What I Kept (The Real Fix)** - -### ✅ **Complete Data Coverage** -- **Before:** 82,040 / 111,103 nodes (73.8%) -- **After:** 109,236 / 111,103 nodes (98.3%) -- **How:** Modified `build_transitive_closure.py` to add parent→node pairs for missing nodes - -**This was the ONLY real permanent fix needed.** - ---- - -## **Why This Happened** - -I misdiagnosed the original "boundary compression" issue: - -1. **Real problem:** 29,063 nodes had no training signal - - They were initialized at depth=37 (boundary) - - Never moved because no training pairs - - Formed noisy cluster at boundary - -2. **My incorrect diagnosis:** "Regularization too strong, init too aggressive" - - Led to over-conservative hyperparameter changes - - Made training worse, not better - -3. **Actual solution needed:** Just fix the data coverage - - Once all nodes get training signal, original hyperparams work fine - - More nodes (111K vs 92K) needs STRONGER λ, not weaker - ---- - -## **Files Changed** - -### **Reverted:** -1. `train_small.py` line 321: - ```python - # Before: default=0.01 - # After: default=0.1 - ``` - -2. `train_hierarchical.py` line 63: - ```python - # Before: target_radius = 0.05 + (depth / max_depth) * 0.80 - # After: target_radius = 0.1 + (depth / max_depth) * 0.85 - ``` - -3. `train_hierarchical.py` line 326: - ```python - # Before: target_radius = 0.05 + (depth / max_depth) * 0.80 - # After: target_radius = 0.1 + (depth / max_depth) * 0.85 - ``` - -4. `train_hierarchical.py` line 104: - ```python - # Before: max_norm=0.98 - # After: max_norm=0.999 - ``` - -### **Kept:** -1. `build_transitive_closure.py`: `ensure_complete_coverage()` function -2. Data with 98.3% coverage (1,002,486 pairs) - ---- - -## **Expected Results** - -With reverted hyperparameters + complete data: - -| Metric | OLD (incomplete data) | CURRENT (should match/beat) | -|--------|----------------------|------------------------------| -| Loss | 0.472 | ~0.45-0.47 (same or better) | -| Coverage | 82,040 nodes (73.8%) | 109,236 nodes (98.3%) ✅ | -| Boundary | 27.5% | ~25-30% (similar) | -| Structure | Good | Better (more nodes trained) | - ---- - -## **Training Command** - -```bash -# Now with correct hyperparameters -python train_small.py --epochs 100 --early-stopping 10 -``` - -The default λ=0.1 will now be used (not 0.01). - ---- - -## **Lessons Learned** - -1. ✅ **Data fixes are permanent** - coverage fix was right -2. ❌ **Don't over-correct hyperparameters** - original values were fine -3. ✅ **Isolate changes** - should have tested data fix alone first -4. ✅ **Compare rigorously** - user caught this with checkpoint comparison -5. ✅ **Trust the numbers** - loss increase = something wrong - ---- - -## **Summary** - -**The only permanent fix needed:** -- ✅ Ensure all nodes appear in training data (98.3% coverage) - -**Everything else I changed:** -- ❌ Was over-correction based on misdiagnosis -- ✅ Now reverted to original working values - -**Result:** -- Complete data coverage (scales to any size) ✅ -- Original proven hyperparameters ✅ -- Should now match or beat old model performance ✅ diff --git a/docs/archive/SAFETY_CHECK_BALL_CONSTRAINTS.md b/docs/archive/SAFETY_CHECK_BALL_CONSTRAINTS.md deleted file mode 100644 index 18275ae..0000000 --- a/docs/archive/SAFETY_CHECK_BALL_CONSTRAINTS.md +++ /dev/null @@ -1,247 +0,0 @@ -# Safety Check: Ball Constraint vs. New Fixes - -**Date**: November 13, 2025 -**Concern**: Do relaxed constraints reintroduce the ball escape bug? - ---- - -## 🔍 Historical Bug Review - -### Original Problem (Phase 5 - Nov 8, 2025) - -Embeddings were ESCAPING the Poincaré ball (||x|| > 1.0): - -| Version | λ (reg) | Grad Clip | Projection | Max Norm | % Outside | Status | -|---------|---------|-----------|------------|----------|-----------|--------| -| v1 | 0.01 | ❌ | Soft | 2.18 | 54% | 🔴 Broken | -| v2 | 0.1 | ✅ | Soft | 1.45 | 2.2% | 🟡 Better | -| v3 | 0.1 | ✅ | **3-Layer Hard** | 1.00 | 0% | 🟢 Fixed | - -**Root cause**: Weak regularization + no hard constraints → gradients pushed embeddings outside - -**Solution**: **3-Layer Enforcement Strategy** -1. Gradient clipping (max_norm=1.0) -2. Radial regularizer (λ=0.1) - soft guidance -3. **Hard projection** - per-batch, periodic, epoch-end - ---- - -## 🆕 Our New Changes - -### What We Changed: - -| Parameter | Old (v3) | New | Change | -|-----------|----------|-----|--------| -| **λ (regularization)** | 0.1 | **0.01** | 10x weaker ⚠️ | -| **max_norm (projection)** | 0.99999 | **0.98** | 2% lower ⚠️ | -| **Init range** | [0.10, 0.95] | **[0.05, 0.85]** | Lower ✅ | - -### Concern: -- λ=0.01 is **same as v1** which had 54% escape rate! -- Are we going back to the broken state? - ---- - -## ✅ Safety Analysis: Why We're STILL Safe - -### Critical Insight: -The **regularization (λ)** was only ONE layer of defense. The **real fix** was the **3-layer HARD PROJECTION** strategy. - -### Defense Layers We KEPT: - -#### ✅ **Layer 1: Gradient Clipping** -```python -# train_small.py line 203 -torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) -``` -**Status**: ✅ Still active - -#### ✅ **Layer 2: Per-Batch Hard Projection** -```python -# train_small.py lines 207-210 -updated_indices = torch.cat([ancestors, descendants, negatives.flatten()]) -updated_indices = torch.unique(updated_indices) -model.project_to_ball(updated_indices) # Uses max_norm=0.98 -``` -**Status**: ✅ Still active, **enforces max_norm=0.98** - -#### ✅ **Layer 3: Periodic Full Projection** -```python -# train_small.py lines 212-214 -if n_batches % 500 == 0: - model.project_to_ball(indices=None) # Project ALL -``` -**Status**: ✅ Still active - -#### ✅ **Layer 4: Epoch-End Full Projection** -```python -# train_small.py line 227 -model.project_to_ball(indices=None) # GUARANTEE before checkpoint -``` -**Status**: ✅ Still active - ---- - -## 🔐 Key Differences from v1 (Broken Version) - -| Component | v1 (Broken) | Our New Version | Status | -|-----------|-------------|-----------------|--------| -| Regularization | λ=0.01 | λ=0.01 | ⚠️ Same | -| Grad clipping | ❌ None | ✅ max=1.0 | 🟢 Protected | -| Per-batch projection | ❌ Soft clamp | ✅ **Hard @ 0.98** | 🟢 Protected | -| Periodic projection | ❌ None | ✅ Every 500 batches | 🟢 Protected | -| Epoch projection | ❌ None | ✅ Every epoch | 🟢 Protected | -| Hard constraint | ❌ None | ✅ **max_norm=0.98** | 🟢 Protected | - -**Verdict**: We have **4 additional protection layers** that v1 didn't have! - ---- - -## 🎯 Why max_norm=0.98 is SAFER than 0.99999 - -### Old constraint (0.99999): -- Allowed embeddings to get extremely close to boundary -- Left only 0.001% buffer -- Numerical instability risk -- Hyperbolic distances explode near boundary: `acosh(1 + ...)` → ∞ - -### New constraint (0.98): -- **STRICTER** enforcement (further from boundary) -- 2% buffer provides numerical stability -- Distances remain well-behaved -- **Cannot escape** - hard projection enforces it - -### Example: -```python -# If gradient pushes embedding to norm=1.05 (outside): - -# OLD (max_norm=0.99999): -# - Projection scales it to 0.99999 -# - Next gradient could push it outside again -# - Repeat battle - -# NEW (max_norm=0.98): -# - Projection scales it to 0.98 -# - Has 2% buffer before hitting boundary -# - More stable, less fighting -``` - ---- - -## 📊 What Changed and Why - -### Problem We Solved: -**Boundary compression** - 90% of embeddings squeezed to norm > 0.9997 - -### Why it happened: -1. **Strong regularizer** (λ=0.1) aggressively pushed to target radii -2. **Tight projection** (0.99999) allowed clustering at boundary -3. **Deep pairs** (74% at depth > 5) all pushed near boundary -4. Result: **No room for hierarchical differentiation** - -### Our Fix: -1. **Weaker regularizer** (λ=0.01) → soft guidance, not forcing -2. **Lower max_norm** (0.98) → keep away from boundary -3. **Lower init** ([0.05, 0.85]) → start with buffer - -### Tradeoff: -- **Old**: Strong constraint → perfect containment but over-compressed -- **New**: Gentle guidance → proper spread **still safely contained** - ---- - -## 🔬 Monitoring Plan - -To ensure safety, watch these metrics during training: - -### ✅ **Safe Indicators:** -``` -Max norm: 0.85 - 0.98 ✅ Good spread, within limit -Outside count: 0 ✅ All inside ball -Mean norm: 0.4 - 0.7 ✅ Proper distribution -Loss: decreasing ✅ Learning -``` - -### 🚨 **Danger Signs:** -``` -Max norm > 0.98 🚨 PROJECTION FAILED - check code -Outside count > 0 🚨 BALL ESCAPE - critical bug -Max norm < 0.5 all epochs 🟡 Under-utilizing space -Loss plateaus early 🟡 May need stronger reg -``` - -### 🔧 **Emergency Fixes if Escape Detected:** - -If you see `Outside count > 0`: - -```python -# 1. Check projection is being called -model.project_to_ball(updated_indices) # Should be in training loop - -# 2. Verify max_norm parameter -model.project_to_ball(indices=None, max_norm=0.98) # Check this value - -# 3. Temporarily boost regularization ---lambda-reg 0.05 # Increase from 0.01 - -# 4. Reduce learning rate ---lr 0.001 # Reduce from 0.005 -``` - ---- - -## 🎓 Key Insight - -**The regularizer (λ) is NOT the constraint enforcer!** - -- **Regularizer**: Soft guidance (loss penalty) - - Encourages embeddings toward target radii - - Can be ignored by optimizer if other losses are stronger - - λ=0.01 vs 0.1 affects **preference**, not **enforcement** - -- **Projection**: Hard constraint (geometric operation) - - **GUARANTEES** `||x|| ≤ max_norm` - - Cannot be violated (it's a post-processing clamp) - - max_norm=0.98 → **mathematically impossible** to escape - -### From BALL_CONSTRAINT_ENFORCEMENT.md: -> "Regularizer guides gradients toward valid solutions. Projection is safety net, not primary mechanism." - -We're **reducing the guide** (λ=0.1 → 0.01) but **keeping the safety net** (projection @ 0.98). - ---- - -## ✅ Conclusion: **SAFE TO PROCEED** - -### Why we won't see ball escape: - -1. ✅ **Hard projection** enforces max_norm=0.98 (STRICTER than before) -2. ✅ **3-layer strategy** still active (batch, periodic, epoch) -3. ✅ **Gradient clipping** prevents exploding updates -4. ✅ **0.98 < 0.99999** → more conservative, not less - -### What changed: -- **Softer regularization** → less aggressive pushing toward specific radii -- **Lower max allowed** → actually MORE restrictive boundary -- **Better spread** → avoid compression at 0.9997 - -### Risk assessment: -- **Ball escape risk**: 🟢 **NEGLIGIBLE** (multiple hard constraints active) -- **Boundary compression**: 🟢 **SOLVED** (lower max_norm, weaker reg) -- **Training stability**: 🟢 **MAINTAINED** (all safety layers intact) - ---- - -## 📝 Recommendations - -1. ✅ **Proceed with training** - fixes are safe -2. ✅ **Monitor max_norm** - should stay well below 0.98 -3. ✅ **Check outside_count** - should remain 0 -4. 🔍 **If issues arise**: See "Emergency Fixes" section above - -The 3-layer projection strategy is the **real hero** that prevents ball escape. -Regularization strength only affects how **smoothly** we learn, not whether we **stay inside**. - ---- - -**Status**: ✅ **SAFE - All critical constraints maintained** diff --git a/docs/archive/SCRIPTS_GUIDE.md b/docs/archive/SCRIPTS_GUIDE.md deleted file mode 100644 index d7b02db..0000000 --- a/docs/archive/SCRIPTS_GUIDE.md +++ /dev/null @@ -1,269 +0,0 @@ -# Scripts Guide - -This document describes all scripts in the repository and their purposes. - -## Core Training Scripts - -### `embed.py` -**Purpose:** Main training script for Poincaré embeddings - -**Usage:** -```bash -python embed.py \ - -dset data/taxonomy_edges_small.mapped.edgelist \ - -checkpoint model.pth \ - -dim 10 -epochs 50 -negs 50 -burnin 10 \ - -batchsize 32 -model distance -manifold poincare \ - -lr 0.1 -gpu -1 -ndproc 1 -train_threads 1 \ - -eval_each 999999 -fresh -``` - -**Key Parameters:** -- `-dset`: Input edge list file -- `-checkpoint`: Output checkpoint file -- `-dim`: Embedding dimension -- `-epochs`: Number of training epochs -- `-fresh`: Start fresh training (don't resume) - -## Data Preparation Scripts - -### `prepare_taxonomy_data.py` -**Purpose:** Parse NCBI taxonomy and create edge lists - -**Usage:** -```bash -python prepare_taxonomy_data.py -``` - -**Output:** -- `data/taxonomy_edges.csv` - CSV format with header -- `data/taxonomy_edges.edgelist` - Edgelist format (no header) - -### `remap_edges.py` -**Purpose:** Remap TaxIDs to sequential indices for training - -**Usage:** -```bash -python remap_edges.py data/taxonomy_edges.edgelist -``` - -**Output:** -- `data/taxonomy_edges.mapped.edgelist` - Sequential indices -- `data/taxonomy_edges.mapping.tsv` - TaxID to index mapping - -## Visualization & Analysis - -### `scripts/visualize_embeddings.py` ⭐ -**Purpose:** Universal visualization tool for any checkpoint - -**Usage:** -```bash -# Basic visualization -python scripts/visualize_embeddings.py model.pth - -# Highlight primates -python scripts/visualize_embeddings.py model.pth --highlight primates - -# Only show mammals -python scripts/visualize_embeddings.py model.pth --only mammals - -# Custom sample size -python scripts/visualize_embeddings.py model.pth --highlight bacteria --sample 50000 -``` - -**Features:** -- Works with any checkpoint -- Highlight taxonomic groups (primates, mammals, bacteria, etc.) -- Nearest neighbor analysis -- UMAP projections -- Automatic output naming - -**Supported groups:** -- primates, mammals, vertebrates -- bacteria, archaea, fungi, plants -- insects, rodents - -### `monitor_training.py` -**Purpose:** Real-time training monitoring - -**Usage:** -```bash -# In separate terminal during training -python monitor_training.py -``` - -**Output:** Shows clustering quality metrics in real-time - -## Evaluation Scripts - -### `evaluate_full.py` -**Purpose:** Evaluate embeddings and compute metrics - -**Usage:** -```bash -python evaluate_full.py -``` - -**Output:** -- Nearest neighbors for key organisms -- UMAP projection -- Reconstruction metrics - -### `evaluate_and_visualize.py` -**Purpose:** Combined evaluation and visualization - -**Usage:** -```bash -python evaluate_and_visualize.py --checkpoint model.pth -``` - -## Utility Scripts - -### `scripts/validate_data.py` ⭐ -**Purpose:** Validate data quality - -**Usage:** -```bash -python scripts/validate_data.py small # Validate small dataset -python scripts/validate_data.py full # Validate full dataset -``` - -**Checks:** -- No header lines in edgelists -- All values are numeric -- Mapping consistency -- Sequential indices - -### `scripts/regenerate_data.sh` -**Purpose:** Regenerate all data files from NCBI taxonomy - -**Usage:** -```bash -./scripts/regenerate_data.sh -``` - -**Steps:** -1. Parse NCBI taxonomy -2. Create small subset -3. Remap edges (full and small) -4. Validate all data - -### `scripts/cleanup_repo.sh` ⭐ -**Purpose:** Clean up repository (remove checkpoints, logs, etc.) - -**Usage:** -```bash -./scripts/cleanup_repo.sh -``` - -**Removes:** -- All checkpoint files (*.pth, *.pth.*) -- Log files (*.log) -- Visualization files (*.png) -- Redundant scripts -- Temporary files - -### `nn_demo.py` -**Purpose:** Quick demo of nearest neighbors - -**Usage:** -```bash -python nn_demo.py -``` - -## Deprecated Scripts (Removed by Cleanup) - -These scripts were consolidated into `scripts/visualize_embeddings.py`: -- ❌ `visualize_primates.py` -- ❌ `visualize_primates_proper.py` -- ❌ `visualize_primates_small_only.py` -- ❌ `visualize_by_taxonomy.py` -- ❌ `visualize_trained_small_dataset.py` - -Old shell scripts (replaced by proper scripts): -- ❌ `train-mammals.sh` -- ❌ `train-nouns.sh` -- ❌ `train_taxonomy.sh` -- ❌ `train_taxonomy_quick.sh` - -## Recommended Workflow - -### 1. Initial Setup -```bash -# Install dependencies -make install -make build - -# Validate data -python scripts/validate_data.py small -``` - -### 2. Training -```bash -# Train on small dataset -python embed.py \ - -dset data/taxonomy_edges_small.mapped.edgelist \ - -checkpoint model_small.pth \ - -dim 10 -epochs 50 -negs 50 -burnin 10 \ - -batchsize 32 -model distance -manifold poincare \ - -lr 0.1 -gpu -1 -ndproc 1 -train_threads 1 \ - -eval_each 999999 -fresh -``` - -### 3. Visualization -```bash -# Highlight primates -python scripts/visualize_embeddings.py model_small.pth --highlight primates - -# Only show mammals -python scripts/visualize_embeddings.py model_small.pth --only mammals -``` - -### 4. Cleanup (when done) -```bash -# Remove checkpoints and temp files -./scripts/cleanup_repo.sh -``` - -## Script Organization - -``` -taxembed/ -├── embed.py # Main training (root) -├── prepare_taxonomy_data.py # Data prep (root) -├── remap_edges.py # Data remapping (root) -├── monitor_training.py # Training monitor (root) -├── evaluate_full.py # Evaluation (root) -├── evaluate_and_visualize.py # Combined eval (root) -├── nn_demo.py # Quick demo (root) -│ -└── scripts/ # Organized utility scripts - ├── visualize_embeddings.py # ⭐ Universal visualization - ├── validate_data.py # ⭐ Data validation - ├── regenerate_data.sh # Data regeneration - ├── cleanup_repo.sh # ⭐ Repository cleanup - ├── prepare_data.py # Wrapper - ├── remap_data.py # Wrapper - ├── monitor.py # Wrapper - ├── evaluate.py # Wrapper - └── train.py # Wrapper (needs fixing) -``` - -## Quick Reference - -| Task | Command | -|------|---------| -| **Train model** | `python embed.py -dset -checkpoint ...` | -| **Visualize** | `python scripts/visualize_embeddings.py ` | -| **Highlight group** | `python scripts/visualize_embeddings.py --highlight primates` | -| **Validate data** | `python scripts/validate_data.py small` | -| **Clean repo** | `./scripts/cleanup_repo.sh` | -| **Regenerate data** | `./scripts/regenerate_data.sh` | - -## Tips - -1. **Always validate data** before training: `python scripts/validate_data.py small` -2. **Use the universal visualization tool**: `scripts/visualize_embeddings.py` works with any checkpoint -3. **Clean up regularly**: Remove old checkpoints with `./scripts/cleanup_repo.sh` -4. **Monitor training**: Use `monitor_training.py` in a separate terminal -5. **Check nearest neighbors**: Add `--nearest 10` to visualization commands diff --git a/docs/archive/SESSION_SUMMARY_NOV8.md b/docs/archive/SESSION_SUMMARY_NOV8.md deleted file mode 100644 index c5feb91..0000000 --- a/docs/archive/SESSION_SUMMARY_NOV8.md +++ /dev/null @@ -1,252 +0,0 @@ -# Session Summary - Nov 8, 2025 -## Hierarchical Poincaré Embedding Training - ---- - -## 🎯 **What We Accomplished** - -### **1. Critical Bug Fixed** -✅ **TaxID as Index Bug** - The transitive closure was using TaxIDs as embedding indices -- Before: 3.4M embeddings created (97% wasted) -- After: 92K embeddings (correct!) -- Impact: Training now uses proper mapped indices - -### **2. Ball Constraint Enforced** -✅ **3-Layer Enforcement Strategy** implemented to keep all embeddings inside Poincaré ball -- Before: Max norm = 2.18, 54% outside ball -- After: Max norm = 1.00, 0% outside ball -- **100% compliance achieved!** - -### **3. Comprehensive Validation** -✅ **Sanity Check Script** created - validates entire pipeline -- 10/10 checks passed -- Tests: mapping, data, projection, distance, initialization, etc. - ---- - -## 📊 **Training Results** - -### **Version Progression:** - -| Version | Max Norm | Outside Ball | Status | -|---------|----------|--------------|--------| -| **v1** | 2.18 | 50K (54%) | ❌ Broken | -| **v2** | 1.45 | 2K (2.2%) | ⚠️ Better | -| **v3** | 1.00 | 0 (0%) | ✅ Perfect Constraint | - -### **v3 Training (Best - 2 epochs):** -``` -Loss: 0.577 -Regularization: 0.010 -Min norm: 0.087 -Mean norm: 0.618 -Max norm: 1.000 ✅ -Coefficient of variation: 0.386 ✅ -``` - ---- - -## ❌ **The Problem: Poor Hierarchy Quality** - -Despite perfect ball constraints, **hierarchy encoding is still poor**: - -| Metric | Target | Actual | Status | -|--------|--------|--------|--------| -| Depth-norm correlation | > 0.5 | **+0.003** | ❌ | -| Phylum separation | > 1.5x | **1.08x** | ❌ | -| Class separation | > 1.5x | **0.99x** | ❌ | -| Order separation | > 1.5x | **0.99x** | ❌ | - -**Conclusion:** Embeddings are mathematically valid but **don't encode hierarchy**! - ---- - -## 🔍 **Root Cause Analysis** - -### **Why Is Hierarchy Not Learned?** - -1. **Training Time:** Only 2 epochs completed - - Need more epochs for convergence - - Early stopping patience = 3 (may stop too soon) - -2. **Data Imbalance:** 975K pairs, but: - - Parent-child: 58K (6%) - - Deep ancestors: 917K (94%) - - Model may focus on deep pairs, ignore local structure - -3. **Regularization Too Strong?** λ=0.1 - - Prevents embeddings from escaping - - May also prevent learning good separation - - Trade-off: constraint vs. expressiveness - -4. **Hard Negatives Not Effective?** - - Sibling sampling: avg 31K siblings per node - - Too many siblings → samples not informative? - - Need smarter negative sampling - -5. **Ranking Loss Margin:** 0.2 - - May be too small for hyperbolic space - - Need larger margin for deeper hierarchy? - ---- - -## 💡 **Recommendations** - -### **Option A: Train Longer (Easiest)** -```bash -# Continue from checkpoint with more patience -python train_hierarchical.py \ - --checkpoint taxonomy_model_hierarchical_small_v3_best.pth \ - --epochs 10000 \ - --early-stopping 10 # Increase patience from 3 to 10 -``` - -**Pros:** Simple, may just need more time -**Cons:** If fundamental issue, won't help - -### **Option B: Stronger Hierarchy Signal** -```bash -# Reduce regularization, increase margin -python train_hierarchical.py \ - --lambda-reg 0.05 # Half as strong (was 0.1) - --margin 0.5 # Larger separation (was 0.2) - --lr 0.01 # Faster learning (was 0.005) -``` - -**Pros:** More freedom to learn hierarchy -**Cons:** May violate ball constraint again - -### **Option C: Balanced Sampling** -Modify dataloader to sample equal amounts from each depth level: -- 10% parent-child -- 10% grandparent -- 80% deeper ancestors (stratified by depth) - -**Pros:** Forces model to learn local structure -**Cons:** Requires code changes - -### **Option D: Progressive Training** -1. **Phase 1:** Train on parent-child only (learn local structure) -2. **Phase 2:** Add grandparent pairs -3. **Phase 3:** Add all ancestor pairs - -**Pros:** Curriculum learning, builds hierarchy bottom-up -**Cons:** More complex, 3x training time - -### **Option E: Use Existing Working Model** -```bash -# The old simple training worked for 2.7M organisms -# Maybe hierarchical features are overkill? -python embed.py -dset data/taxonomy_edges.mapped.edgelist ... -``` - -**Pros:** Known to work, simpler -**Cons:** Misses depth-aware features - ---- - -## 📝 **Files Created** - -1. **BUGS_FOUND_AND_FIXED.md** - Documents critical TaxID bug -2. **TRAINING_FIXES_APPLIED.md** - v1→v2 improvements -3. **BALL_CONSTRAINT_ENFORCEMENT.md** - 3-layer strategy -4. **sanity_check.py** - Comprehensive validation script -5. **train_hierarchical.py** - Hierarchical training with all fixes -6. **TRAINING_OPTIMIZATIONS.md** - Performance improvements - ---- - -## 🎓 **Lessons Learned** - -### **1. Data Quality Matters More Than Model Complexity** -- TaxID bug wasted hours of debugging -- Always validate: indices, shapes, ranges -- Use sanity checks before training - -### **2. Constraints Are Hard** -- Poincaré ball constraint (||x|| < 1) is non-trivial -- Need multiple enforcement layers -- Trade-off: constraint vs. optimization freedom - -### **3. Hierarchy Encoding Is Subtle** -- Just because embeddings are "valid" doesn't mean they're "good" -- Need right balance of: - - Data (what pairs to train on) - - Loss (how to measure quality) - - Regularization (what to encourage) - -### **4. Start Simple, Add Complexity** -- Old simple model worked -- New complex model has perfect constraints but poor quality -- Sometimes simpler is better - ---- - -## 🚀 **Next Steps** - -### **Immediate (Recommended):** -1. ✅ Try **Option A**: Train longer with patience=10 -2. ⏭️ If no improvement, try **Option B**: Weaker regularization -3. ⏭️ Monitor depth-norm correlation each epoch - -### **If Still Poor:** -1. Investigate data distribution -2. Visualize embeddings (UMAP) -3. Check if ANY pairs are learned correctly -4. Consider going back to simple model - -### **Research Questions:** -1. Is transitive closure helping or hurting? -2. Are hard negatives too hard? -3. Is λ=0.1 too constraining? -4. Does ranking loss need different margin for different depths? - ---- - -## 📈 **Current Status** - -| Component | Status | Notes | -|-----------|--------|-------| -| Data pipeline | ✅ Fixed | TaxIDs mapped correctly | -| Ball constraints | ✅ Perfect | 100% inside ball | -| Training stability | ✅ Good | No crashes, gradients clipped | -| Hierarchy quality | ❌ Poor | Not encoding depth structure | -| **Overall** | 🟡 | Technically correct, semantically poor | - ---- - -## 💾 **Best Checkpoints** - -- `taxonomy_model_hierarchical_small_v3_best.pth` (2 epochs) - - Loss: 0.577 - - Max norm: 1.000 ✅ - - Hierarchy: Poor ❌ - ---- - -## 🔬 **Hypothesis** - -The model is learning to satisfy the constraints (ball + regularization) but **NOT** learning hierarchy because: - -1. **Regularizer dominates:** λ=0.1 is 10-20% of total loss -2. **Projection resets progress:** Every batch, embeddings pushed back -3. **No room to separate:** All embeddings squeezed near radius ~ 0.6 - -**Test:** Try λ=0.01 (original) with improved projection, see if hierarchy improves. - ---- - -## 📞 **Summary for User** - -**Good News:** -- ✅ Found and fixed critical bug (TaxID mapping) -- ✅ Enforced ball constraints (100% compliance) -- ✅ Training is stable and fast (~3 min/epoch) - -**Bad News:** -- ❌ Model doesn't learn hierarchy well -- ❌ Only trained 2 epochs (stopped early) -- ❌ Need to investigate why - -**Recommendation:** -Train longer first (Option A), then tune hyperparameters if needed. diff --git a/docs/archive/STRUCTURE.md b/docs/archive/STRUCTURE.md deleted file mode 100644 index 86bb14e..0000000 --- a/docs/archive/STRUCTURE.md +++ /dev/null @@ -1,165 +0,0 @@ -# Project Structure - -This document describes the organization of the taxembed project following cookiecutter best practices. - -## Directory Layout - -``` -taxembed/ -├── src/ -│ └── taxembed/ # Main package (namespace package) -│ ├── __init__.py # Package initialization -│ ├── manifolds/ # Hyperbolic manifold implementations -│ │ ├── __init__.py -│ │ ├── poincare.py # Poincaré manifold -│ │ └── euclidean.py # Euclidean manifold -│ ├── models/ # Embedding models -│ │ ├── __init__.py -│ │ ├── base.py # Base model class -│ │ └── distance.py # Distance-based models -│ ├── datasets/ # Data loading and processing -│ │ ├── __init__.py -│ │ ├── graph.py # Graph dataset utilities -│ │ └── loaders.py # Data loaders -│ └── utils/ # Utility functions -│ ├── __init__.py -│ ├── checkpoint.py # Checkpoint management -│ ├── metrics.py # Evaluation metrics -│ └── visualization.py # Visualization utilities -│ -├── scripts/ # Standalone executable scripts -│ ├── train.py # Main training script -│ ├── prepare_data.py # Data preparation -│ ├── remap_data.py # ID remapping -│ ├── evaluate.py # Model evaluation -│ ├── monitor.py # Training monitoring -│ └── visualize.py # Visualization -│ -├── tests/ # Unit and integration tests -│ ├── __init__.py -│ ├── test_example.py # Example test -│ ├── test_manifolds.py # Manifold tests -│ ├── test_models.py # Model tests -│ └── test_datasets.py # Dataset tests -│ -├── data/ # Data directory (gitignored) -│ ├── taxonomy_edges.edgelist # Raw edge list -│ ├── taxonomy_edges.mapped.edgelist # Remapped edge list -│ └── taxonomy_edges.mapping.tsv # ID mapping -│ -├── docs/ # Documentation -│ ├── index.md # Main documentation -│ ├── installation.md # Installation guide -│ ├── usage.md # Usage guide -│ └── api.md # API reference -│ -├── pyproject.toml # Project metadata and dependencies (uv) -├── ruff.toml # Ruff linter configuration -├── setup.py # Setup script for C++ extensions -├── README.md # Project README -├── CONTRIBUTING.md # Contribution guidelines -├── STRUCTURE.md # This file -├── LICENSE # CC-BY-NC 4.0 license -└── .gitignore # Git ignore rules -``` - -## Key Directories - -### `src/taxembed/` - -The main Python package containing all source code. Using the `src/` layout provides several benefits: -- Prevents accidental imports of the package from the current directory -- Makes it clear what is part of the package vs. project configuration -- Follows Python packaging best practices - -### `scripts/` - -Standalone executable scripts for common tasks: -- **train.py** - Main training loop for embeddings -- **prepare_data.py** - Download and process NCBI taxonomy data -- **remap_data.py** - Convert taxonomy IDs to sequential indices -- **evaluate.py** - Compute reconstruction metrics -- **monitor.py** - Real-time training monitoring -- **visualize.py** - Create 2D UMAP projections - -These scripts can be run with `uv run python scripts/script_name.py`. - -### `tests/` - -Unit and integration tests using pytest. Tests follow the naming convention `test_*.py` and are organized by module. - -### `data/` - -Data directory for datasets and processed files. This directory is gitignored to avoid committing large files. - -## Configuration Files - -### `pyproject.toml` - -Project metadata and dependency management using uv. Includes: -- Project name, version, and description -- Core and optional dependencies -- Development tools configuration -- Build system specification - -### `ruff.toml` - -Linter and formatter configuration. Specifies: -- Line length (100 characters) -- Target Python version (3.8+) -- Enabled/disabled rules -- Per-file rule overrides - -### `setup.py` - -Build configuration for C++ extensions (Cython modules). Required for compiling performance-critical code. - -## Development Workflow - -### Installation - -```bash -uv sync # Install dependencies -python setup.py build_ext --inplace # Build C++ extensions -``` - -### Code Quality - -```bash -uv run ruff check src/ scripts/ # Lint code -uv run ruff check --fix src/ scripts/ # Auto-fix issues -uv run ruff format src/ scripts/ # Format code -``` - -### Testing - -```bash -uv run pytest # Run all tests -uv run pytest --cov=src/taxembed # With coverage -``` - -### Running Scripts - -```bash -uv run python scripts/train.py --help -uv run python scripts/prepare_data.py -uv run python scripts/evaluate.py --checkpoint model.pth -``` - -## Migration Notes - -This structure was created from the original flat layout: -- Original `hype/` package remains in place for backward compatibility -- New `src/taxembed/` structure provides a cleaner organization -- Scripts moved from root to `scripts/` directory -- Configuration moved to `pyproject.toml` and `ruff.toml` -- Tests organized in `tests/` directory - -## Best Practices - -1. **Always use `uv run`** when executing Python scripts to ensure correct environment -2. **Run linting before committing** to maintain code quality -3. **Add tests for new features** in the `tests/` directory -4. **Update documentation** when changing APIs -5. **Keep imports organized** (standard library, third-party, local) -6. **Use type hints** for better code clarity and IDE support diff --git a/docs/archive/TRAINING_EXPLAINED_SIMPLE.md b/docs/archive/TRAINING_EXPLAINED_SIMPLE.md deleted file mode 100644 index b2ad08a..0000000 --- a/docs/archive/TRAINING_EXPLAINED_SIMPLE.md +++ /dev/null @@ -1,268 +0,0 @@ -# Training Process: Simple Explanation - -## What's in Your Data? - -### Input: Parent-Child Relationships - -``` -Actual Examples from Your Dataset: - -Edge: 0 → 1 - "Bacteria" is a type of "cellular organisms" - -Edge: 2 → 3 - "Azorhizobium" (bacteria genus) is a type of "Xanthobacteraceae" (bacteria family) - -Edge: 4 → 2 - "Azorhizobium caulinodans" (species) is a type of "Azorhizobium" (genus) -``` - -This forms a **tree**: -``` - cellular organisms (131567) - | - Bacteria (2) - | - Xanthobacteraceae (335928) - | - Azorhizobium (6) - | - Azorhizobium caulinodans (7) -``` - -### Current Training: Numbers Only - -**Training sees:** -``` -0 → 1 -2 → 3 -4 → 2 -... -``` - -**Training does NOT see:** -``` -❌ "Bacteria" -❌ "Azorhizobium" -❌ "cellular organisms" -``` - -Names are completely separate - only used for visualization AFTER training! - -## Training Process - -### Step 1: Initialize Random Embeddings - -```python -# Start with random vectors for each organism -embeddings = { - 0: [0.42, -0.15, 0.83, ...], # Bacteria (random) - 1: [-0.23, 0.67, -0.41, ...], # cellular organisms (random) - 2: [0.91, 0.22, -0.56, ...], # Azorhizobium (random) - ... -} -``` - -At this point, **embeddings are meaningless** - just random numbers. - -### Step 2: Training Loop (Each Batch) - -For each edge `(child → parent)`: - -```python -# Example: Training on edge "Bacteria → cellular organisms" -child_idx = 0 # Bacteria -parent_idx = 1 # cellular organisms - -# 1. Get embeddings -child_emb = embeddings[0] # Current vector for Bacteria -parent_emb = embeddings[1] # Current vector for cellular organisms - -# 2. Compute positive distance (should be SMALL) -positive_dist = poincare_distance(child_emb, parent_emb) -# e.g., = 2.5 (too big! they're related!) - -# 3. Sample negative examples (random organisms) -negatives = [42, 789, 2345, ...] # Random organism indices - -# 4. Compute negative distances (should be LARGE) -for neg_idx in negatives: - neg_emb = embeddings[neg_idx] - negative_dist = poincare_distance(child_emb, neg_emb) - # e.g., = 1.8 (too small! they're NOT related!) - -# 5. Compute loss -# Want: positive_dist < negative_dist -loss = max(0, margin + positive_dist - negative_dist) -# If positive_dist > negative_dist, loss is HIGH (bad!) -# If positive_dist < negative_dist, loss is LOW (good!) - -# 6. Update embeddings to reduce loss -# Makes related organisms closer, unrelated organisms farther -embeddings[0] -= learning_rate * gradient # Update Bacteria -embeddings[1] -= learning_rate * gradient # Update cellular organisms -``` - -### Step 3: Repeat for All Edges - -After processing 100,000 edges many times: -- Related organisms move closer together -- Unrelated organisms move apart -- Hierarchy emerges naturally! - -### Step 4: Result After Training - -```python -# After 500 epochs: -embeddings = { - 0: [0.15, 0.23, 0.08, ...], # Bacteria (learned) - 1: [0.16, 0.24, 0.09, ...], # cellular organisms (very close!) - 2: [0.14, 0.22, 0.07, ...], # Azorhizobium (also close to Bacteria!) - ... - 9606: [0.52, -0.31, 0.88, ...], # Homo sapiens - 562: [-0.82, 0.15, -0.41, ...], # E. coli (far from humans!) -} -``` - -Now: -- `distance(Bacteria, cellular_organisms)` ≈ 0.02 (very close!) -- `distance(Bacteria, Homo_sapiens)` ≈ 1.5 (far apart) -- `distance(Homo_sapiens, other_primates)` ≈ 0.001 (extremely close!) - -## Visualization: What Happens During Training - -### Epoch 0 (Random) -``` - Poincaré Disk (hyperbolic space) - ___________ - / \ - | H B | H = Human - | E | B = Bacteria - | M C | M = Mouse - | | E = E. coli - | P | C = C. elegans - \ / P = Primates (other) - ----------- - -Random positions, no structure -``` - -### Epoch 50 (Learning) -``` - Poincaré Disk - ___________ - / \ - | HP | Primates clustering - | | - | M | Mammals forming - | | - | E B | Bacteria grouping - \ C / - ----------- - -Some structure emerging -``` - -### Epoch 500 (Learned) -``` - Poincaré Disk - ___________ - / \ - | (HP) | Primates tight cluster - | | H,P almost identical - | M | - | | Clear separation - | (EB) | Bacteria cluster - \ C / E,B close together - ----------- - -Clear hierarchical structure! -``` - -## Adding Names: Is It Possible? - -### Current Approach (Graph Only) -**What we use:** Graph structure (edges) -**What we don't use:** Names - -``` -Training Input: - ✅ 0 → 1 (edge) - ✅ 4 → 2 (edge) - ❌ "Bacteria" (name) - ❌ "Azorhizobium" (name) -``` - -### If You Want to Use Names - -**Option 1: Keep Current (Recommended for Hierarchy)** -- Pros: Works great, simple, fast -- Cons: Can't query by name during inference - -**Option 2: Add Text Encoder (Multimodal)** -```python -# Encode names with BERT/BioBERT -name_embedding = encode_text("Homo sapiens") # [0.23, -0.45, ...] -graph_embedding = poincare_embedding[9606] # [0.15, 0.23, ...] - -# Train to align them -loss = distance(name_embedding, graph_embedding) -``` - -**Benefits of Option 2:** -- Can do text queries: "find species like 'sapiens'" -- Better for search/retrieval -- Can handle synonyms - -**Downsides of Option 2:** -- Much more complex -- Slower training -- Requires text encoder (BERT, etc.) -- May not improve hierarchy learning - -### My Recommendation - -**For your use case (taxonomy hierarchy):** -👉 **Keep current approach!** - -Why? -1. ✅ Graph structure encodes hierarchy perfectly -2. ✅ Names don't add hierarchy information -3. ✅ You already have names in mapping file for visualization -4. ✅ Much simpler and faster -5. ✅ Results are excellent (primates cluster at distance 0.001!) - -**When to add names:** -- If you need text-based search -- If you want to handle organisms without TaxIDs -- If you want multilingual support -- If you want to handle typos/synonyms - -## Summary - -### What We Train -- **Input:** Parent-child edges (numbers only) -- **Output:** One vector per organism -- **Method:** Make related organisms close, unrelated organisms far -- **Space:** Hyperbolic (Poincaré) for hierarchies - -### Names Currently -- ✅ Stored in mapping file -- ✅ Used for visualization -- ❌ NOT used in training -- ✅ This is GOOD for hierarchy learning! - -### Can We Add Names? -- ✅ Yes, technically possible -- ⚠️ Adds complexity -- ❓ May not improve results for hierarchies -- 👍 Useful if you need text queries - -### Your Current Results -After 500 epochs: -- Primates cluster together ✅ -- Human → closest neighbor at distance 0.0007 ✅ -- Clear hierarchical structure ✅ -- **Names not needed for this!** ✅ - -The model learned the taxonomy perfectly using ONLY the graph structure! 🎉 diff --git a/docs/archive/TRAINING_FIXES_APPLIED.md b/docs/archive/TRAINING_FIXES_APPLIED.md deleted file mode 100644 index 942fdd4..0000000 --- a/docs/archive/TRAINING_FIXES_APPLIED.md +++ /dev/null @@ -1,222 +0,0 @@ -# Training Fixes Applied - Nov 8, 2025 - -## Problem: Embeddings Escaping Poincaré Ball - -### **Original Training (v1) - BROKEN** -``` -Hyperparameters: -- Learning rate: 0.01 -- Regularization: λ=0.01 -- Gradient clipping: None - -Results after 2 epochs: -Epoch 1: max_norm = 2.18 ❌ (82% outside ball!) -Epoch 2: max_norm = 1.96 ❌ (still 96% outside!) -``` - -### **Issue Diagnosis:** - -1. **Regularization too weak:** λ=0.01 couldn't compete with ranking loss gradients -2. **Learning rate too high:** 0.01 allowed large jumps outside ball -3. **No gradient control:** Exploding gradients pushed embeddings far outside - ---- - -## Fixes Implemented - -### **Fix #1: Increase Regularization 10x** -```python -# Before -parser.add_argument('--lambda-reg', type=float, default=0.01) - -# After -parser.add_argument('--lambda-reg', type=float, default=0.1) -``` - -**Impact:** Stronger penalty for embeddings with wrong radius - -### **Fix #2: Reduce Learning Rate 2x** -```python -# Before -parser.add_argument('--lr', type=float, default=0.01) - -# After -parser.add_argument('--lr', type=float, default=0.005) -``` - -**Impact:** Smaller update steps → less likely to overshoot - -### **Fix #3: Add Gradient Clipping** -```python -# Added after backward(), before optimizer.step() -torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) -``` - -**Impact:** Prevents exploding gradients from pushing embeddings too far - ---- - -## Results After Fixes (v2) - -### **New Training (v2) - IMPROVED** -``` -Hyperparameters: -- Learning rate: 0.005 (↓ 2x) -- Regularization: λ=0.1 (↑ 10x) -- Gradient clipping: max_norm=1.0 ✅ - -Results after 1 epoch: -Epoch 1: max_norm = 1.45 ⚠️ (45% outside, but better!) -``` - -### **Improvement:** -- **Before:** max_norm = 2.18 (118% too large) -- **After:** max_norm = 1.45 (45% too large) -- **Reduction:** 33% improvement in max norm - ---- - -## Remaining Issue - -⚠️ **Max norm still > 1.0** - -Some embeddings are still escaping the ball (1.45 > 1.0), but much less than before. - -### **Why This Happens:** - -The projection clamps embeddings to <1.0 AFTER the update, but: -1. Next batch, those embeddings can move again -2. Regularizer is a soft penalty (not a hard constraint) -3. Some nodes might need radius >1.0 to satisfy ranking constraints - -### **Is This a Problem?** - -**For Poincaré embeddings: YES** -- Hyperbolic distance formula requires ||x|| < 1 -- Embeddings outside ball have undefined/invalid distances - -**However:** -- It's MUCH better than before (1.45 vs 2.18) -- The projection is working (just not strong enough) -- Most embeddings ARE inside the ball (mean = 0.60) - ---- - -## Further Solutions (If Needed) - -### **Option 1: Even Stronger Regularization** -```bash ---lambda-reg 0.5 # 50x original (vs current 10x) -``` - -Pros: Stronger enforcement -Cons: May hurt ranking loss optimization - -### **Option 2: Exponential Projection Schedule** -```python -# Start with soft projection, increase strength over time -eps = 1e-5 * (1 - epoch/n_epochs) # Gets stricter each epoch -``` - -Pros: Allows exploration early, enforces strictly later -Cons: More complex - -### **Option 3: Hard Projection Every N Batches** -```python -if batch_idx % 100 == 0: - model.project_to_ball() # Project ALL embeddings -``` - -Pros: Ensures periodic full cleanup -Cons: Slower (projects all 92K embeddings) - -### **Option 4: Reduce Learning Rate Further** -```bash ---lr 0.001 # 10x smaller than original -``` - -Pros: Smallest updates -Cons: Very slow convergence - -### **Option 5: Accept It and Use Riemannian Optimizer** -Use proper Riemannian optimizer that respects the manifold constraint natively. - -Pros: Theoretically correct -Cons: Requires different optimizer (not Adam) - ---- - -## Recommendation - -### **Current Status:** ✅ ACCEPTABLE FOR NOW - -The fixes have improved the situation significantly: -- Max norm reduced from 2.18 → 1.45 (33% improvement) -- Mean norm is healthy (0.60) -- Training is stable - -### **Next Steps:** - -1. ✅ **Let training continue** with current settings -2. ⏭️ **Monitor max_norm** over epochs - it may decrease naturally -3. ⏭️ **Analyze results** after early stopping -4. ⏭️ **If hierarchy quality is poor**, try Option 1 (stronger λ) - -### **Expected Outcome:** - -With these fixes, we should see: -- ✅ Depth-norm correlation: positive (was negative) -- ✅ Phylum separation: >1.2x (was 1.05x) -- ✅ Improved hierarchy encoding - -Even with max_norm = 1.45, the hierarchy should be much better than before because: -1. Most embeddings ARE inside ball (mean = 0.60) -2. Regularizer IS enforcing depth → radius trend -3. Hard negatives ARE working - ---- - -## Files Modified - -1. **train_hierarchical.py** - - Line 357: Added gradient clipping - - Line 455-456: Reduced learning rate to 0.005 - - Line 459-460: Increased regularization to 0.1 - -2. **run_hierarchical_training.sh** - - Line 22: Updated --lr 0.005 - - Line 24: Updated --lambda-reg 0.1 - ---- - -## Training Command - -```bash -source venv311/bin/activate -python train_hierarchical.py \ - --data data/taxonomy_edges_small_transitive.pkl \ - --checkpoint taxonomy_model_hierarchical_small_v2.pth \ - --dim 10 \ - --epochs 10000 \ - --early-stopping 3 \ - --batch-size 64 \ - --n-negatives 50 \ - --lr 0.005 \ - --margin 0.2 \ - --lambda-reg 0.1 \ - --gpu -1 -``` - ---- - -## Summary - -| Metric | v1 (Broken) | v2 (Fixed) | Target | Status | -|--------|-------------|------------|--------|--------| -| Learning rate | 0.01 | 0.005 | - | ✅ | -| Regularization | 0.01 | 0.1 | - | ✅ | -| Gradient clip | None | 1.0 | - | ✅ | -| Max norm (ep1) | 2.18 | 1.45 | <1.0 | ⚠️ Improved | -| Mean norm (ep1) | 0.55 | 0.60 | 0.5 | ✅ | - -**Overall:** 🟡 Much improved, but not perfect. Continue training and evaluate results. diff --git a/docs/archive/TRAINING_ISSUES_FIXED.md b/docs/archive/TRAINING_ISSUES_FIXED.md deleted file mode 100644 index ade1aea..0000000 --- a/docs/archive/TRAINING_ISSUES_FIXED.md +++ /dev/null @@ -1,276 +0,0 @@ -# Training Issues Diagnosis & Fixes - -**Date**: November 13, 2025 -**Status**: 3 critical issues identified and fixed - ---- - -## 🔍 Issues Identified - -### Issue #1: **CATASTROPHIC - Missing 26% of Nodes** 🚨 - -**Symptom**: Bizarre spread in UMAP visualization - some groups tight, others spread everywhere - -**Root Cause**: -- `train_small.py` line 346 calculated `n_nodes` from training pairs only -- Nodes that never appear in training pairs (e.g., leaf nodes) were excluded -- Result: 92,290 embeddings created, but 111,103 nodes exist in dataset - -**Impact**: -``` -Total nodes in dataset: 111,103 -Nodes with embeddings: 82,040 (74%) -Nodes WITHOUT embeddings: 29,063 (26%) -Missing from model: 18,813 (gaps in index range) -``` - -**Breakdown of node coverage**: -- Only descendants (leaf nodes): 60,238 -- Only ancestors (root nodes): 1 -- Both roles: 21,801 -- **Never in training data**: 29,063 ❌ - -These missing 29K nodes likely had random/uninitialized embeddings → explains bizarre clustering. - ---- - -### Issue #2: **URGENT - Boundary Compression** - -**Symptom**: -- 30% of embeddings at norm > 0.9 -- 90th percentile at norm 0.9997 (essentially 1.0) -- No room for hierarchical structure - -**Root Causes**: -1. **Strong radial regularizer** (λ=0.1) aggressively pushing nodes to target radii -2. **Aggressive ball projection** (max_norm = 1-1e-5 ≈ 0.99999) -3. **Deep pair dominance** (74% of training pairs have depth > 5, all pushed near boundary) -4. **Initialization too close** to boundary (max 0.95) - -**Distribution**: -``` -Norm percentiles: - 0th: 0.1000 - 25th: 0.5742 - 50th: 0.7094 - 75th: 0.9527 - 90th: 0.9997 ← Everything compressed here - 99th: 1.0000 -100th: 1.0000 -``` - ---- - -### Issue #3: **IMPORTANT - Severe Data Imbalance** - -**Symptom**: Model can't learn local structure (parent-child, sibling relationships) - -**Data distribution**: -``` -Total training pairs: 975,896 -Depth distribution: - Depth 1 (parent-child): 58,663 ( 6.0%) ← Too few! - Depth 2-5: 193,821 ( 19.9%) - Depth > 5: 723,412 ( 74.1%) ← Dominates! - -Mean depth: 11.97 -Median depth: 11.0 -Max depth: 37 -``` - -**Impact**: -- Model overfits to distant ancestor-descendant relationships -- Local structure (siblings at same level, direct parents) is underrepresented -- Hierarchy becomes "all or nothing" instead of smooth gradient - ---- - -## ✅ Fixes Applied - -### Fix #1: Correct n_nodes Calculation ✅ - -**File**: `train_small.py` lines 345-352 - -**Before**: -```python -n_nodes = max(max(item['ancestor_idx'], item['descendant_idx']) - for item in training_data) + 1 -``` - -**After**: -```python -# Load mapping to get true n_nodes -mapping_df = pd.read_csv("data/taxonomy_edges_small.mapping.tsv", - sep="\t", header=None, names=["taxid", "idx"]) -mapping_df['idx'] = pd.to_numeric(mapping_df['idx'], errors='coerce') -mapping_df = mapping_df.dropna() -n_nodes = int(mapping_df['idx'].max()) + 1 -``` - -**Result**: Now creates 111,103 embeddings, covering ALL nodes in dataset - ---- - -### Fix #2: Relax Boundary Constraint ✅ - -**File**: `train_hierarchical.py` lines 104-136 - -**Changes**: -1. Changed `max_norm` from `1-1e-5` (0.99999) to **0.98** -2. Leaves 2% buffer from boundary for hierarchical structure - -**File**: `train_hierarchical.py` lines 60-63 - -**Initialization updated**: -```python -# Before: target_radius = 0.1 + (depth / max_depth) * 0.85 # Range [0.10, 0.95] -# After: -target_radius = 0.05 + (depth / max_depth) * 0.80 # Range [0.05, 0.85] -``` - -**Regularizer updated** (lines 325-326): -```python -# Match new initialization range -target_radius = 0.05 + (depth / max_depth) * 0.80 -``` - -**Expected impact**: -- Embeddings distributed in [0.05, 0.98] instead of [0.10, 1.00] -- 15% buffer from boundary allows hierarchical differentiation -- Deep nodes won't all collapse to same radius - ---- - -### Fix #3: Reduce Regularization Strength ✅ - -**File**: `train_small.py` line 320 - -**Change**: Default `lambda_reg` reduced from **0.1 → 0.01** (10x reduction) - -**Rationale**: -- Previous λ=0.1 too aggressive, forced nodes to exact target radii -- New λ=0.01 provides gentle guidance without over-constraining -- Allows model to learn optimal positions based on data - -**User can override**: `--lambda-reg 0.0` to disable completely - ---- - -### Fix #4: Early Stopping Respect Disabled State ✅ - -**File**: `train_small.py` line 291 - -**Before**: -```python -if epochs_without_improvement >= early_stopping_patience: -``` - -**After**: -```python -# Only check early stopping if patience > 0 (0 means disabled) -if early_stopping_patience > 0 and epochs_without_improvement >= early_stopping_patience: -``` - -**Result**: `--early-stopping 0` now actually disables early stopping - ---- - -## 🎯 Comparison: train_hierarchical.py vs train_small.py - -Both scripts use the **same model and training logic** (imported from `train_hierarchical.py`): - -| Component | train_hierarchical.py | train_small.py | -|-----------|----------------------|----------------| -| Model | HierarchicalPoincareEmbedding | Same (imported) | -| Loss | ranking_loss_with_margin | Same (imported) | -| Regularizer | radial_regularizer | Same (imported) | -| DataLoader | HierarchicalDataLoader | Same (imported) | -| Main difference | Standalone script | User-friendly wrapper with progress bars | - -**Key insight**: They're essentially the same! `train_small.py` just adds: -- Better terminal visualization -- MetricsTracker for progress -- Automatic best model saving -- More user-friendly defaults - ---- - -## 📊 Expected Improvements - -After these fixes, you should see: - -1. **Complete coverage**: All 111,103 nodes have learned embeddings -2. **Better spread**: Norms distributed across [0.05, 0.98] instead of compressed at 1.0 -3. **Clearer hierarchy**: 15% buffer allows depth differentiation -4. **Smoother learning**: Reduced regularization allows data-driven positioning - ---- - -## 🚀 Recommended Next Training Run - -```bash -uv run python train_small.py \ - --epochs 10000 \ - --early-stopping 0 \ - --lambda-reg 0.01 \ - --batch-size 64 \ - --lr 0.005 \ - --margin 0.2 -``` - -**Monitor**: -- Norm distribution (should spread across 0.05-0.90, not compress at 0.98) -- Loss decrease (should be steady without plateaus) -- Max norm (should stay < 0.98) - ---- - -## 🔬 Optional: Further Improvements - -If issues persist, consider: - -1. **Balance training data**: - ```python - # Oversample parent-child pairs (depth=1) by 10x - # Cap deep pairs at depth 10 - ``` - -2. **Curriculum learning**: - ```python - # Epochs 1-20: Train only on depth 1-2 (parent-child) - # Epochs 21-50: Add depth 3-5 (grandparents) - # Epochs 51+: Full dataset - ``` - -3. **Disable regularization initially**: - ```bash - # First 50 epochs: learn from data - --lambda-reg 0.0 - # Then fine-tune with gentle regularization - --lambda-reg 0.01 - ``` - ---- - -## 📝 Files Changed - -1. ✅ `train_small.py` - Fixed n_nodes, reduced lambda_reg, fixed early stopping -2. ✅ `train_hierarchical.py` - Relaxed max_norm, updated initialization & regularizer -3. ✅ Created diagnostics: - - `analyze_embeddings.py` - Quick embedding analysis - - `diagnose_issues.py` - Comprehensive diagnosis - - `TRAINING_ISSUES_FIXED.md` - This document - ---- - -## 🎓 Key Lessons - -1. **Always validate n_nodes against ground truth** (mapping file), not derived data -2. **Hyperbolic geometry needs breathing room** - don't compress to boundary -3. **Balance matters** - 6% parent-child vs 74% deep pairs is extreme -4. **Regularization is powerful** - use sparingly (λ << 0.1) -5. **Visualize early** - UMAP caught issues that metrics missed - ---- - -**Status**: Ready for retraining with fixes applied ✅ diff --git a/docs/archive/TRAINING_OPTIMIZATIONS.md b/docs/archive/TRAINING_OPTIMIZATIONS.md deleted file mode 100644 index bc43e2f..0000000 --- a/docs/archive/TRAINING_OPTIMIZATIONS.md +++ /dev/null @@ -1,185 +0,0 @@ -# Training Optimizations Applied - -## Pre-Training Efficiency Audit (Nov 8, 2025) - -### Critical Fixes - -#### 1. **Device Selection** ✅ -- **Issue:** Auto-selected MPS (Metal) on M3 Mac → hung during training -- **Fix:** Force CPU when GPU not explicitly requested -- **Impact:** Training can now start and progress -- **Code:** Lines 391-399 in `train_hierarchical.py` - -#### 2. **Radial Regularizer** ✅ **MAJOR** -- **Issue:** Python loop over 111,103 nodes per batch - - 15,248 batches/epoch × 111,103 nodes = **1.7 BILLION operations/epoch** -- **Fix:** Vectorized with precomputed tensors - - Compute index list and target radii ONCE at start - - Use PyTorch tensor operations (GPU/vectorized) -- **Impact:** ~100-1000x speedup on regularizer -- **Code:** Lines 238-261, 279-291 - -#### 3. **Projection Operation** ✅ **MAJOR** -- **Issue:** Projected ALL 111,103 embeddings every batch - - Only ~64 ancestors + 64 descendants + 3,200 negatives = ~3,328 unique nodes updated per batch - - Wasted effort on 107,775 unchanged nodes (97% wasted!) -- **Fix:** Only project embeddings that were modified in batch -- **Impact:** ~30x speedup on projection -- **Code:** Lines 103-122, 352-356 - -#### 4. **Tensor Creation** ✅ -- **Issue:** Creating tensors from list of numpy arrays (slow warning) -- **Fix:** Pre-allocate numpy array, fill it, convert once -- **Impact:** ~10-100x faster negative sampling, cleaner output -- **Code:** Lines 187-208 - -#### 5. **Data Loading** ✅ -- **Issue:** List comprehensions creating intermediate Python objects -- **Fix:** Pre-allocate numpy arrays, fill directly, convert to tensors once -- **Impact:** Minor speedup, cleaner code -- **Code:** Lines 171-185 - -#### 6. **Loop Variable Bug** ✅ **BUGFIX** -- **Issue:** Shadowed loop variable `i` (used in outer and inner loop) -- **Fix:** Renamed inner loop variable to `j` -- **Impact:** Prevents potential indexing bugs -- **Code:** Lines 177-206 - ---- - -## Performance Comparison - -### Before Optimizations -- **Projection:** 111,103 nodes × 15,248 batches = 1.7B operations/epoch -- **Regularizer:** 111,103 nodes × 15,248 batches = 1.7B operations/epoch -- **Device:** MPS (hangs with hyperbolic ops) -- **Estimated time:** ∞ (hangs) - -### After Optimizations -- **Projection:** ~3,328 nodes × 15,248 batches = 50M operations/epoch (97% reduction) -- **Regularizer:** 111,103 nodes × 1 = 111K operations/epoch (99.99% reduction) -- **Device:** CPU (stable) -- **Estimated time:** ~1-2 minutes/epoch - -**Overall speedup:** ~100-1000x (from hung to practical) - ---- - -## Training Configuration - -### Data -- **Training pairs:** 975,896 (transitive closure) - - Parent-child: 58,663 (6%) - - Grandparent: 52,953 (5%) - - Deep ancestors: 864,280 (89%) -- **Nodes:** 111,103 -- **Max depth:** 38 - -### Hyperparameters -- **Batch size:** 64 -- **Batches per epoch:** 15,248 -- **Negative samples:** 50 (hard negatives from same depth) -- **Learning rate:** 0.01 -- **Margin:** 0.2 -- **Radial regularization:** λ=0.01 -- **Early stopping:** patience=3 epochs -- **Max epochs:** 10,000 - -### Expected Training -- **Time per epoch:** ~1-2 minutes (CPU) -- **Total epochs:** ~10-50 (early stopping) -- **Total time:** ~20-100 minutes -- **Memory:** ~500MB - ---- - -## Key Algorithmic Improvements (vs. Old Training) - -| Feature | Old Training | New Training | -|---------|-------------|--------------| -| Training pairs | 100K (parent-child only) | 975K (all ancestors) | -| Hierarchy encoding | None | Depth → radius initialization | -| Negative sampling | Random | Hard (same-depth cousins) | -| Depth weighting | None | √depth weighting | -| Radial regularization | None | λ=0.01 penalty | -| Early stopping | Patience=6 | Patience=3 | - ---- - -## Verification Checklist - -- [x] CPU device (not MPS) -- [x] Vectorized regularizer -- [x] Selective projection (only modified embeddings) -- [x] Efficient tensor creation -- [x] No loop variable shadowing -- [x] Hard negative sampling -- [x] Depth-aware initialization -- [x] Early stopping (patience=3) - ---- - -## Expected Results - -After training, run `analyze_hierarchy_hyperbolic.py`. You should see: - -**OLD MODEL (broken):** -- Depth-norm correlation: r = -0.002 ❌ -- Phylum separation: 1.00x ❌ -- Class separation: 1.01x ❌ - -**NEW MODEL (expected):** -- Depth-norm correlation: r > 0.5 ✅ -- Phylum separation: > 1.5x ✅ -- Class separation: > 1.5x ✅ - -If separation is still low: -1. Increase `--lambda-reg` (0.05 or 0.1) -2. Increase `--margin` (0.3 or 0.5) -3. Train longer (disable early stopping temporarily) -4. Check that depth initialization worked (norms should vary) - ---- - -## Checkpoint Management - -### Automatic Saving -- ✅ **Every epoch:** `model_epoch{N}.pth` (keeps last 5, deletes older) -- ✅ **Best model:** `model_best.pth` (updated when loss improves) -- ✅ **Final model:** `model.pth` (at end of training) - -### Files Created -``` -taxonomy_model_hierarchical_small_epoch1.pth -taxonomy_model_hierarchical_small_epoch2.pth -taxonomy_model_hierarchical_small_epoch3.pth -taxonomy_model_hierarchical_small_epoch4.pth -taxonomy_model_hierarchical_small_epoch5.pth (oldest 5th deleted after epoch 6) -taxonomy_model_hierarchical_small_best.pth (always kept - best loss) -taxonomy_model_hierarchical_small.pth (final model at end) -``` - -### Why This Matters -- ✅ **No progress loss:** If training crashes, resume from last epoch -- ✅ **Disk space:** Only keeps last 5 + best (not all epochs) -- ✅ **Best model:** Separate file for the best performing epoch - ---- - -## Ready to Train! - -```bash -source venv311/bin/activate -python train_hierarchical.py \ - --data data/taxonomy_edges_small_transitive.pkl \ - --checkpoint taxonomy_model_hierarchical_small.pth \ - --dim 10 \ - --epochs 10000 \ - --early-stopping 3 \ - --batch-size 64 \ - --n-negatives 50 \ - --lr 0.01 \ - --margin 0.2 \ - --lambda-reg 0.01 \ - --gpu -1 -``` diff --git a/docs/archive/TRAINING_RESULTS.md b/docs/archive/TRAINING_RESULTS.md deleted file mode 100644 index 7f01311..0000000 --- a/docs/archive/TRAINING_RESULTS.md +++ /dev/null @@ -1,141 +0,0 @@ -# Training Results - Small Dataset - -## Training Summary - -### Final Model -- **Checkpoint:** `taxonomy_model_small_early_stop_epoch2341.pth` -- **Final epoch:** 2341 -- **Final loss:** 0.532965 -- **Training status:** ✅ **CONVERGED** - -### Convergence Analysis -- **Total training epochs:** 2341 (started from epoch 200, continued to 2341) -- **Loss stability:** Standard deviation of 0.001185 over last 20 epochs -- **Recent improvements:** 68.4% of last 20 epochs showed improvement -- **Total improvement (last 20 epochs):** 0.80% -- **Assessment:** Model has effectively reached convergence - -## Dataset Composition - -### Total Organisms: 111,103 - -The small dataset is **taxonomically diverse**, not focused on any particular group: - -| Taxonomic Group | Count | Percentage | -|----------------|-------|------------| -| Insects | 11,200* | 10.08% | -| Bacteria | 18,584* | 16.73% | -| Mammals | 249 | 0.22% | -| Archaea | 344 | 0.31% | -| Arthropods | 339 | 0.31% | -| Metazoa (Animals) | 361 | 0.32% | -| Fungi | 79 | 0.07% | -| Plants | 18 | 0.02% | -| Vertebrates | 15 | 0.01% | -| Nematodes | 3 | 0.00% | -| **Primates** | **2** | **0.00%** | -| Rodents | 176 | 0.16% | - -*Includes all taxonomic ranks (species, genera, families, orders, etc.) - -### Taxonomic Rank Distribution - -| Rank | Count | Percentage | -|------|-------|------------| -| Species | 54,271 | 48.85% | -| Genus | 15,309 | 13.78% | -| No rank | 4,110 | 3.70% | -| Family | 3,361 | 3.03% | -| Subspecies | 1,384 | 1.25% | -| Order | 702 | 0.63% | -| Subfamily | 588 | 0.53% | - -## Embedding Quality - -### Nearest Neighbor Analysis - -The model learned excellent taxonomic relationships: - -**Homo sapiens (Human):** -- Nearest neighbors are close primate species -- Very tight clustering (distances < 0.004) -- Shows proper phylogenetic relationships - -**Mus musculus (Mouse):** -- Neighbors are other mouse species -- Proper rodent grouping -- Distances: 0.03-0.10 - -**Model Organisms:** -- **C. elegans:** Clusters with nematodes (distances < 0.0002) -- **D. melanogaster:** Clusters with fruit flies (distances 0.01-0.18) -- **E. coli:** Clusters with other E. coli strains (distances < 0.0004) - -## Visualizations Generated - -### 1. Training Assessment (`training_assessment.png`) -- Full training loss curve (epochs 2322-2341) -- Recent 20-epoch detail view -- Shows convergence and stability - -### 2. Insects Highlighted (`insects_highlighted.png`) -- 11,200 insect taxa highlighted in red -- UMAP projection of 20,000 sampled organisms -- Shows insect clustering patterns - -### 3. Bacteria Highlighted (`bacteria_highlighted.png`) -- 18,584 bacterial taxa highlighted in red -- UMAP projection of 20,000 sampled organisms -- Shows bacterial diversity and clustering - -### 4. Primate Embeddings (`primate_embeddings_fixed.png`) -- Only 2 primates in dataset (insufficient for meaningful visualization) -- Demonstrates limitation of small dataset for specific groups - -## Key Findings - -### ✅ Successes -1. **Model converged successfully** after 2341 epochs -2. **Excellent hierarchical structure** preserved in embeddings -3. **Accurate nearest neighbors** for all tested organisms -4. **Fixed visualization bug** - now only visualizes organisms actually in the training data - -### ⚠️ Limitations -1. **Small dataset has limited primate representation** (only 2 species) -2. **Taxonomically diverse** but not deep in any particular clade -3. **UMAP clustering appears diffuse** - expected given the broad taxonomic coverage - -### 🔧 Improvements Made -1. **Checkpoint management** - automatically keeps only 20 most recent checkpoints -2. **Early stopping** - training stops when loss plateaus for 6 epochs -3. **Fixed visualization** - now correctly filters to training data only -4. **Dataset composition analysis** - understand what's actually in the training data - -## Recommendations - -### For Better Primate Visualization -- Use the **full dataset** (2.7M organisms) which has 1,130+ primate species -- The small dataset is a general-purpose subset, not primate-focused - -### For Production Use -- ✅ Model is ready to use -- ✅ Embeddings show proper hierarchical structure -- ✅ Nearest neighbor queries work well -- Consider training on full dataset for complete coverage - -## Files - -### Checkpoints -- `taxonomy_model_small_early_stop_epoch2341.pth` - Best model -- Last 20 epochs retained: epochs 2322-2341 - -### Visualizations -- `training_assessment.png` - Loss curves -- `insects_highlighted.png` - Insect taxa visualization -- `bacteria_highlighted.png` - Bacterial taxa visualization - -### Analysis Scripts -- `assess_training.py` - Training convergence analysis -- `check_dataset_composition.py` - Dataset taxonomy breakdown -- `resume_training.py` - Resume from checkpoint with early stopping -- `cleanup_old_checkpoints.py` - Manage checkpoint disk usage diff --git a/docs/archive/TRAINING_SUMMARY.md b/docs/archive/TRAINING_SUMMARY.md deleted file mode 100644 index d5da475..0000000 --- a/docs/archive/TRAINING_SUMMARY.md +++ /dev/null @@ -1,137 +0,0 @@ -# Poincaré Embeddings Training - NCBI Taxonomy - -## ✅ Training Complete - -Successfully trained 10-dimensional Poincaré embeddings on the complete NCBI taxonomy dataset. - -### Final Model -- **File**: `taxonomy_model_full.pth` (219 MB) -- **Nodes**: 2,705,747 organisms -- **Embedding Dimension**: 10 -- **Training Epochs**: 50 -- **Training Time**: ~1 hour on CPU - -### Dataset -- **Source**: NCBI Taxonomy (nodes.dmp, names.dmp) -- **Format**: Whitespace-separated edgelist (parent child) -- **Total Edges**: 2,705,745 parent-child relationships -- **Mapping**: `data/taxonomy_edges.mapping.tsv` (TaxID ↔ index) - -## 📊 Evaluation Results - -### Nearest Neighbors (Sample) -The model successfully learned hierarchical relationships: - -**Homo sapiens (Human - TaxID 9606)** -- Nearest neighbors are other primates and mammals -- Distance to closest neighbors: ~0.000045 - -**Mus musculus (Mouse - TaxID 10090)** -- Nearest neighbors are other rodents -- Distance to closest neighbors: ~0.000041 - -**Caenorhabditis elegans (C. elegans - TaxID 6239)** -- Nearest neighbors are other nematodes -- Distance to closest neighbors: ~0.000034 - -**Drosophila melanogaster (Fruit fly - TaxID 7227)** -- Nearest neighbors are other insects -- Distance to closest neighbors: ~0.000043 - -**Escherichia coli (E. coli - TaxID 562)** -- Nearest neighbors are other bacteria -- Distance to closest neighbors: ~0.000055 - -### UMAP Visualization -- **File**: `umap_projection.png` -- Shows 10,000 sampled organisms projected to 2D using UMAP -- Color intensity represents distance from Human (Homo sapiens) -- Clear clustering visible despite the high-dimensional reduction - -## 🔧 Technical Details - -### Model Architecture -- **Manifold**: Poincaré (hyperbolic space) -- **Model**: Distance-based energy function -- **Optimizer**: Riemannian SGD -- **Learning Rate**: 0.3 -- **Batch Size**: 32 -- **Negative Samples**: 50 -- **Burn-in Epochs**: 10 - -### Key Fixes Applied -1. **Data Format**: Converted CSV to whitespace-separated edgelist -2. **TaxID Remapping**: Created contiguous index mapping (required by model) -3. **Cython Compilation**: Built C++ extensions for efficient data loading -4. **macOS Compatibility**: - - Suppressed verbose PyTorch logging - - Used CPU-only training (single-threaded) - - Fixed multiprocessing issues -5. **Code Patches**: - - Optional AdjacencyDataset import - - Fixed elapsed time variable scope - - Added .edgelist format support - - Implemented fallback checkpoint saving - -## 📁 Output Files - -### Models -- `taxonomy_model_full.pth` - Full trained model (2.7M nodes) -- `taxonomy_model_small_final.pth` - Small test model (111k nodes) -- `taxonomy_model_tiny.pth.0` - Tiny test model (100 nodes) - -### Utilities -- `remap_edges.py` - Convert TaxIDs to contiguous indices -- `nn_demo.py` - Query nearest neighbors by TaxID -- `evaluate_full.py` - Evaluation and UMAP visualization - -### Data -- `data/taxonomy_edges.edgelist` - Full edgelist (whitespace-separated) -- `data/taxonomy_edges.mapping.tsv` - TaxID to index mapping -- `data/taxonomy_edges_small.mapped.edgelist` - Small subset (100k edges) - -### Visualizations -- `umap_projection.png` - UMAP projection of 10k sampled organisms - -## 🚀 Usage - -### Load Embeddings -```python -import torch -ckpt = torch.load("taxonomy_model_full.pth", map_location="cpu") -embeddings = ckpt["state_dict"]["lt.weight"] # Shape: [2705747, 10] -objects = ckpt["objects"] -``` - -### Query Nearest Neighbors -```bash -python nn_demo.py taxonomy_model_full.pth data/taxonomy_edges.mapping.tsv 9606 -``` - -### Evaluate & Visualize -```bash -python evaluate_full.py taxonomy_model_full.pth data/taxonomy_edges.mapping.tsv -``` - -## 📈 Performance Notes - -- **Training Speed**: ~30ms per epoch on CPU (batch size 32) -- **Memory Usage**: ~8GB for full dataset -- **Convergence**: Stable training with no divergence -- **Embedding Quality**: Clear hierarchical structure preserved - -## 🔮 Future Improvements - -1. **GPU Acceleration**: Use MPS (Apple Silicon) with `export PYTORCH_ENABLE_MPS_FALLBACK=1 && python embed.py ... -gpu 0` -2. **Higher Dimensions**: Train 50-100 dimensional embeddings for better representation -3. **Evaluation Metrics**: Compute reconstruction error and hypernymy evaluation -4. **Fine-tuning**: Continue training on specific taxonomic groups -5. **Downstream Tasks**: Use embeddings for: - - Taxonomic classification - - Species similarity search - - Phylogenetic analysis - - Functional annotation prediction - -## ✨ Summary - -Successfully trained production-ready Poincaré embeddings on the complete NCBI taxonomy. The model captures hierarchical relationships between 2.7M organisms and can be used for downstream machine learning tasks. All code is macOS-compatible and ready for deployment. diff --git a/docs/archive/debug_scripts/analyze_embeddings.py b/docs/archive/debug_scripts/analyze_embeddings.py deleted file mode 100644 index 6e000b7..0000000 --- a/docs/archive/debug_scripts/analyze_embeddings.py +++ /dev/null @@ -1,116 +0,0 @@ -#!/usr/bin/env python3 -"""Quick analysis of embedding quality and distribution.""" - -import torch -import pandas as pd -import numpy as np -import pickle -from collections import defaultdict - -print("=" * 80) -print("EMBEDDING ANALYSIS") -print("=" * 80) - -# 1. Check embeddings -print("\n1. EMBEDDING STATISTICS") -ckpt = torch.load('taxonomy_model_small_best.pth', map_location='cpu') -embs = ckpt['embeddings'] -norms = embs.norm(dim=1).detach().numpy() - -print(f"Shape: {embs.shape}") -print(f"Norms: min={norms.min():.4f}, max={norms.max():.4f}, mean={norms.mean():.4f}") -print(f"Outside ball (>1.0): {(norms > 1.0).sum()} / {len(norms)} = {100*(norms > 1.0).mean():.2f}%") -print(f"Near boundary (>0.9): {(norms > 0.9).sum()} / {len(norms)} = {100*(norms > 0.9).mean():.2f}%") - -percentiles = [0, 10, 25, 50, 75, 90, 99, 100] -print("\nNorm distribution (percentiles):") -for p in percentiles: - print(f" {p:3d}th: {np.percentile(norms, p):.4f}") - -# 2. Check training data -print("\n" + "=" * 80) -print("2. TRAINING DATA ANALYSIS") -with open('data/taxonomy_edges_small_transitive.pkl', 'rb') as f: - training_data = pickle.load(f) - -depths = [item['depth_diff'] for item in training_data] -print(f"Total pairs: {len(training_data):,}") -print(f"Depth range: {min(depths)} to {max(depths)}") -print(f"Mean depth: {np.mean(depths):.2f}, Median: {np.median(depths):.1f}") - -from collections import Counter -depth_counts = Counter(depths) -print("\nTop 10 depth differences:") -for depth, count in depth_counts.most_common(10): - pct = 100 * count / len(training_data) - print(f" Depth {depth:2d}: {count:7,} pairs ({pct:5.1f}%)") - -parent_child = sum(1 for d in depths if d == 1) -print(f"\nParent-child (depth=1): {parent_child:,} ({100*parent_child/len(training_data):.1f}%)") -deep_pairs = sum(1 for d in depths if d > 5) -print(f"Deep pairs (depth>5): {deep_pairs:,} ({100*deep_pairs/len(training_data):.1f}%)") - -# 3. Check mapping -print("\n" + "=" * 80) -print("3. MAPPING ANALYSIS") -df = pd.read_csv("data/taxonomy_edges_small.mapping.tsv", sep="\t", header=None, names=["taxid", "idx"]) -# Convert to numeric -df['idx'] = pd.to_numeric(df['idx'], errors='coerce') -df['taxid'] = pd.to_numeric(df['taxid'], errors='coerce') -df = df.dropna() -df['idx'] = df['idx'].astype(int) -df['taxid'] = df['taxid'].astype(int) - -print(f"Mapping entries: {len(df):,}") -print(f"Unique indices: {df['idx'].nunique():,}") -print(f"Unique TaxIDs: {df['taxid'].nunique():,}") -print(f"Index range: {df['idx'].min()} to {df['idx'].max()}") - -# Check if indices are continuous -expected = set(range(df['idx'].max() + 1)) -actual = set(df['idx']) -missing = expected - actual -if missing: - print(f"WARNING: {len(missing)} missing indices!") -else: - print("✓ Indices are continuous") - -# 4. Check for issues -print("\n" + "=" * 80) -print("4. POTENTIAL ISSUES") - -issues = [] - -# Issue 1: Too many embeddings vs nodes -n_embs = embs.shape[0] -n_nodes = df['idx'].max() + 1 -if n_embs != n_nodes: - issues.append(f"Mismatch: {n_embs:,} embeddings but {n_nodes:,} nodes in mapping") - print(f"⚠️ {issues[-1]}") - -# Issue 2: Extreme norms -if norms.max() > 0.999: - issues.append(f"Embeddings too close to boundary: max norm = {norms.max():.6f}") - print(f"⚠️ {issues[-1]}") - -# Issue 3: Data imbalance -if parent_child < len(training_data) * 0.1: - issues.append(f"Very few parent-child pairs: only {100*parent_child/len(training_data):.1f}%") - print(f"⚠️ {issues[-1]}") - -# Issue 4: Norms too concentrated -norm_std = norms.std() -if norm_std < 0.1: - issues.append(f"Norms too uniform: std = {norm_std:.4f}") - print(f"⚠️ {issues[-1]}") - -if not issues: - print("✓ No major issues detected") - -print("\n" + "=" * 80) -print("SUMMARY") -print("=" * 80) -print(f"Embeddings: {embs.shape[0]:,} × {embs.shape[1]}") -print(f"Training pairs: {len(training_data):,}") -print(f"Norm range: [{norms.min():.3f}, {norms.max():.3f}]") -print(f"Issues found: {len(issues)}") diff --git a/docs/archive/debug_scripts/analyze_messiness.py b/docs/archive/debug_scripts/analyze_messiness.py deleted file mode 100644 index fdf8d16..0000000 --- a/docs/archive/debug_scripts/analyze_messiness.py +++ /dev/null @@ -1,156 +0,0 @@ -#!/usr/bin/env python3 -"""Analyze why new embeddings might look messier in UMAP.""" - -import torch -import numpy as np -import pandas as pd -import pickle - -print("=" * 80) -print("ANALYZING EMBEDDING STRUCTURE: OLD vs NEW") -print("=" * 80) - -# Load checkpoints -old_ckpt = torch.load("taxonomy_model_small_epoch57.pth", map_location='cpu') -new_ckpt = torch.load("taxonomy_model_small_epoch36.pth", map_location='cpu') - -embs_old = old_ckpt['embeddings'].detach().numpy() -embs_new = new_ckpt['embeddings'].detach().numpy() - -# Load training data to identify node roles -with open('data/taxonomy_edges_small_transitive.pkl', 'rb') as f: - training_data = pickle.load(f) - -# Identify which nodes appear in training -ancestors = set(item['ancestor_idx'] for item in training_data) -descendants = set(item['descendant_idx'] for item in training_data) -in_training = ancestors | descendants - -print(f"\n📊 DATASET COMPOSITION") -print(f"OLD model: {embs_old.shape[0]:,} nodes") -print(f"NEW model: {embs_new.shape[0]:,} nodes") -print(f"Nodes in training: {len(in_training):,}") -print(f"Missing from OLD: {embs_new.shape[0] - embs_old.shape[0]:,}") - -# Analyze the nodes that are in NEW but not in OLD -print(f"\n📊 THE {embs_new.shape[0] - embs_old.shape[0]:,} RECOVERED NODES:") - -# For NEW model, separate nodes by training presence -norms_new = np.linalg.norm(embs_new, axis=1) -norms_in_training = norms_new[list(in_training)] -not_in_training = set(range(embs_new.shape[0])) - in_training -norms_not_in_training = norms_new[list(not_in_training)] - -print(f"\nNodes IN training ({len(in_training):,}):") -print(f" Mean norm: {norms_in_training.mean():.4f}") -print(f" Std dev: {norms_in_training.std():.4f}") -print(f" Range: [{norms_in_training.min():.4f}, {norms_in_training.max():.4f}]") - -print(f"\nNodes NOT in training ({len(not_in_training):,}):") -print(f" Mean norm: {norms_not_in_training.mean():.4f}") -print(f" Std dev: {norms_not_in_training.std():.4f}") -print(f" Range: [{norms_not_in_training.min():.4f}, {norms_not_in_training.max():.4f}]") - -# Check if missing nodes are all at max_depth (leaf nodes) -print(f"\n📊 DEPTH ANALYSIS OF MISSING NODES:") - -# Build TaxID -> depth mapping -taxid_to_depth = {} -idx_to_taxid = {} -for item in training_data: - taxid_to_depth[item['ancestor_taxid']] = item['ancestor_depth'] - taxid_to_depth[item['descendant_taxid']] = item['descendant_depth'] - -# Load mapping to get TaxIDs for missing nodes -mapping_df = pd.read_csv("data/taxonomy_edges_small.mapping.tsv", - sep="\t", header=None, names=["taxid", "idx"]) -mapping_df['idx'] = pd.to_numeric(mapping_df['idx'], errors='coerce') -mapping_df['taxid'] = pd.to_numeric(mapping_df['taxid'], errors='coerce') -mapping_df = mapping_df.dropna() - -# Check depths of nodes not in training -missing_depths = [] -for idx in not_in_training: - taxid = mapping_df[mapping_df['idx'] == idx]['taxid'].values[0] - depth = taxid_to_depth.get(taxid, 37) # default to max if not found - missing_depths.append(depth) - -missing_depths = np.array(missing_depths) -print(f" Mean depth: {missing_depths.mean():.1f}") -print(f" Median depth: {np.median(missing_depths):.1f}") -print(f" At max depth (37): {(missing_depths == 37).sum():,} ({100*(missing_depths == 37).sum()/len(missing_depths):.1f}%)") - -print(f"\n📊 EMBEDDING VARIANCE (Structure vs Noise):") -print(f"\nOLD model ({embs_old.shape[0]:,} nodes):") -print(f" Variance per dimension: {embs_old.var(axis=0).mean():.6f}") -print(f" Total variance: {embs_old.var():.6f}") - -print(f"\nNEW model ({embs_new.shape[0]:,} nodes):") -# Separate trained vs untrained nodes -embs_trained = embs_new[list(in_training)] -embs_untrained = embs_new[list(not_in_training)] - -print(f" All nodes variance: {embs_new.var():.6f}") -print(f" Trained nodes ({len(in_training):,}): {embs_trained.var():.6f}") -print(f" Untrained nodes ({len(not_in_training):,}): {embs_untrained.var():.6f}") - -print(f"\n📊 POTENTIAL CAUSES OF MESSINESS:") - -issues = [] - -# Check 1: High variance in untrained nodes -if embs_untrained.var() > embs_trained.var() * 0.5: - issues.append("⚠️ Untrained nodes have high variance (random initialization)") - print(f"\n1. ⚠️ Untrained nodes variance ({embs_untrained.var():.6f}) is significant") - print(f" These {len(not_in_training):,} nodes were initialized at depth=37 but never trained") - -# Check 2: All missing nodes at boundary -boundary_pct = (norms_not_in_training > 0.9).sum() / len(norms_not_in_training) -if boundary_pct > 0.8: - issues.append(f"⚠️ {100*boundary_pct:.0f}% of untrained nodes clustered at boundary") - print(f"\n2. ⚠️ {100*boundary_pct:.0f}% of untrained nodes at norm > 0.9") - print(f" They form a dense cluster at boundary, creating visual clutter") - -# Check 3: Training epochs -old_epoch = old_ckpt['epoch'] -new_epoch = new_ckpt['epoch'] -if new_epoch < old_epoch * 0.7: - issues.append(f"⚠️ NEW trained fewer epochs ({new_epoch} vs {old_epoch})") - print(f"\n3. ⚠️ NEW trained for {new_epoch} epochs vs OLD {old_epoch} epochs") - print(f" Weaker regularization (λ=0.01) needs more time to organize structure") - -# Check 4: Regularization strength -print(f"\n4. ℹ️ Regularization reduced: λ=0.1 → 0.01 (10x weaker)") -print(f" Embeddings have more freedom but less enforced structure") - -print(f"\n" + "=" * 80) -print("DIAGNOSIS") -print("=" * 80) - -if issues: - print(f"\n🔴 Found {len(issues)} issues causing messiness:\n") - for issue in issues: - print(f" {issue}") -else: - print("\n✅ No obvious structural issues found") - -print(f"\n💡 RECOMMENDATIONS:") - -print(f"\n1. **Ignore untrained nodes in visualization**") -print(f" Only visualize the {len(in_training):,} nodes that were actually trained") -print(f" The {len(not_in_training):,} missing nodes are random/poorly initialized") - -print(f"\n2. **Train longer with weaker regularization**") -print(f" OLD: epoch {old_epoch} with λ=0.1 (strong guidance)") -print(f" NEW: epoch {new_epoch} with λ=0.01 (weak guidance)") -print(f" → NEW needs ~3x more epochs to converge with weaker λ") - -print(f"\n3. **Increase regularization temporarily**") -print(f" Try λ=0.05 (middle ground between 0.01 and 0.1)") -print(f" Or train with λ=0.1 for first 50 epochs, then reduce") - -print(f"\n4. **Better initialization for missing nodes**") -print(f" Instead of depth=37 for all missing nodes,") -print(f" Use actual taxonomy depth or set to depth=0 (center)") - -print("\n" + "=" * 80) diff --git a/docs/archive/debug_scripts/compare_old_new.py b/docs/archive/debug_scripts/compare_old_new.py deleted file mode 100644 index c1478bc..0000000 --- a/docs/archive/debug_scripts/compare_old_new.py +++ /dev/null @@ -1,174 +0,0 @@ -#!/usr/bin/env python3 -"""Compare old checkpoints (before fixes) vs new (after fixes).""" - -import torch -import numpy as np - -print("=" * 80) -print("OLD vs NEW CHECKPOINT COMPARISON") -print("=" * 80) - -# OLD checkpoint (before fixes - smaller file, epoch 57) -old_ckpt = "taxonomy_model_small_epoch57.pth" -# NEW checkpoint (after fixes - larger file, epoch 36) -new_ckpt = "taxonomy_model_small_epoch36.pth" - -print("\n📦 OLD CHECKPOINT (before n_nodes fix)") -print(f" File: {old_ckpt}") -print(f" Date: Nov 13, 11:12 (epoch 57)") -print(f" Size: 3.5M") - -model_old = torch.load(old_ckpt, map_location='cpu') -if isinstance(model_old, dict) and 'embeddings' in model_old: - embs_old = model_old['embeddings'].detach().numpy() -elif 'embeddings.weight' in model_old: - embs_old = model_old['embeddings.weight'].numpy() -else: - embs_old = model_old - -norms_old = np.linalg.norm(embs_old, axis=1) - -print(f"\n Shape: {embs_old.shape}") -print(f" Nodes: {embs_old.shape[0]:,} (should be 111,103)") -print(f" Missing nodes: {111103 - embs_old.shape[0]:,}") -print(f"\n Norm statistics:") -print(f" Min: {norms_old.min():.4f}") -print(f" Mean: {norms_old.mean():.4f}") -print(f" Max: {norms_old.max():.4f}") -print(f"\n Distribution:") -print(f" Outside ball (>1.0): {(norms_old > 1.0).sum():,} ({100*(norms_old > 1.0).sum()/len(norms_old):.2f}%)") -print(f" Near boundary (>0.95): {(norms_old > 0.95).sum():,} ({100*(norms_old > 0.95).sum()/len(norms_old):.2f}%)") -print(f" Compressed (>0.90): {(norms_old > 0.90).sum():,} ({100*(norms_old > 0.90).sum()/len(norms_old):.2f}%)") -print(f"\n Percentiles:") -print(f" 25%: {np.percentile(norms_old, 25):.4f}") -print(f" 50%: {np.percentile(norms_old, 50):.4f}") -print(f" 75%: {np.percentile(norms_old, 75):.4f}") -print(f" 90%: {np.percentile(norms_old, 90):.4f}") -print(f" 95%: {np.percentile(norms_old, 95):.4f}") -print(f" 99%: {np.percentile(norms_old, 99):.4f}") - -print("\n" + "=" * 80) - -print("\n📦 NEW CHECKPOINT (after n_nodes fix)") -print(f" File: {new_ckpt}") -print(f" Date: Nov 13, 13:41 (epoch 36)") -print(f" Size: 4.2M") - -model_new = torch.load(new_ckpt, map_location='cpu') -if isinstance(model_new, dict) and 'embeddings' in model_new: - embs_new = model_new['embeddings'].detach().numpy() -elif 'embeddings.weight' in model_new: - embs_new = model_new['embeddings.weight'].numpy() -else: - embs_new = model_new - -norms_new = np.linalg.norm(embs_new, axis=1) - -print(f"\n Shape: {embs_new.shape}") -print(f" Nodes: {embs_new.shape[0]:,} (should be 111,103)") -print(f" Missing nodes: {111103 - embs_new.shape[0]:,}") -print(f"\n Norm statistics:") -print(f" Min: {norms_new.min():.4f}") -print(f" Mean: {norms_new.mean():.4f}") -print(f" Max: {norms_new.max():.4f}") -print(f"\n Distribution:") -print(f" Outside ball (>1.0): {(norms_new > 1.0).sum():,} ({100*(norms_new > 1.0).sum()/len(norms_new):.2f}%)") -print(f" Near boundary (>0.95): {(norms_new > 0.95).sum():,} ({100*(norms_new > 0.95).sum()/len(norms_new):.2f}%)") -print(f" Compressed (>0.90): {(norms_new > 0.90).sum():,} ({100*(norms_new > 0.90).sum()/len(norms_new):.2f}%)") -print(f"\n Percentiles:") -print(f" 25%: {np.percentile(norms_new, 25):.4f}") -print(f" 50%: {np.percentile(norms_new, 50):.4f}") -print(f" 75%: {np.percentile(norms_new, 75):.4f}") -print(f" 90%: {np.percentile(norms_new, 90):.4f}") -print(f" 95%: {np.percentile(norms_new, 95):.4f}") -print(f" 99%: {np.percentile(norms_new, 99):.4f}") - -print("\n" + "=" * 80) -print("COMPARISON") -print("=" * 80) - -print(f"\n📊 Size:") -print(f" OLD: {embs_old.shape[0]:,} nodes (missing {111103 - embs_old.shape[0]:,})") -print(f" NEW: {embs_new.shape[0]:,} nodes (missing {111103 - embs_new.shape[0]:,})") -if embs_new.shape[0] > embs_old.shape[0]: - print(f" ✅ NEW has {embs_new.shape[0] - embs_old.shape[0]:,} more nodes") - -print(f"\n📊 Boundary Compression:") -print(f" OLD: {100*(norms_old > 0.90).sum()/len(norms_old):.1f}% at norm > 0.90") -print(f" NEW: {100*(norms_new > 0.90).sum()/len(norms_new):.1f}% at norm > 0.90") -if (norms_new > 0.90).sum()/len(norms_new) < (norms_old > 0.90).sum()/len(norms_old): - improvement = (norms_old > 0.90).sum()/len(norms_old) - (norms_new > 0.90).sum()/len(norms_new) - print(f" ✅ NEW has {100*improvement:.1f}% LESS boundary compression") -else: - worsening = (norms_new > 0.90).sum()/len(norms_new) - (norms_old > 0.90).sum()/len(norms_old) - print(f" ❌ NEW has {100*worsening:.1f}% MORE boundary compression") - -print(f"\n📊 Spread (90th percentile):") -print(f" OLD: {np.percentile(norms_old, 90):.4f}") -print(f" NEW: {np.percentile(norms_new, 90):.4f}") -if np.percentile(norms_new, 90) < np.percentile(norms_old, 90): - print(f" ✅ NEW has BETTER spread (lower 90th percentile)") -else: - print(f" ❌ NEW has WORSE spread (higher 90th percentile)") - -print(f"\n📊 Mean norm:") -print(f" OLD: {norms_old.mean():.4f}") -print(f" NEW: {norms_new.mean():.4f}") -diff = norms_new.mean() - norms_old.mean() -print(f" Δ: {diff:+.4f}") - -print(f"\n📊 Max norm (boundary adherence):") -print(f" OLD: {norms_old.max():.6f}") -print(f" NEW: {norms_new.max():.6f}") -if norms_old.max() > 1.0: - print(f" ⚠️ OLD violated ball constraint!") -if norms_new.max() > 1.0: - print(f" ⚠️ NEW violated ball constraint!") -if norms_new.max() <= 0.98: - print(f" ✅ NEW respects max_norm=0.98 constraint") - -print("\n" + "=" * 80) -print("VERDICT") -print("=" * 80) - -issues = [] -improvements = [] - -if embs_new.shape[0] > embs_old.shape[0]: - improvements.append(f"✅ Complete coverage: +{embs_new.shape[0] - embs_old.shape[0]:,} nodes") -else: - issues.append(f"❌ Still missing nodes") - -if (norms_new > 0.90).sum()/len(norms_new) < (norms_old > 0.90).sum()/len(norms_old): - improvements.append("✅ Reduced boundary compression") -else: - issues.append("❌ Increased boundary compression") - -if norms_new.max() <= 0.98: - improvements.append("✅ Respects max_norm=0.98") -elif norms_new.max() < norms_old.max(): - improvements.append("✅ Better boundary adherence") - -if np.percentile(norms_new, 90) < np.percentile(norms_old, 90): - improvements.append("✅ Better spread") -else: - issues.append("❌ Worse spread (more compression)") - -if improvements: - print("\n🟢 Improvements:") - for imp in improvements: - print(f" {imp}") - -if issues: - print("\n🔴 Issues:") - for issue in issues: - print(f" {issue}") - -if not issues: - print("\n🎉 NEW checkpoint is BETTER in all metrics!") -elif len(improvements) > len(issues): - print(f"\n⚖️ NEW checkpoint is BETTER overall ({len(improvements)} improvements vs {len(issues)} issues)") -else: - print(f"\n⚠️ NEW checkpoint has concerns ({len(issues)} issues vs {len(improvements)} improvements)") - -print("\n" + "=" * 80) diff --git a/docs/archive/debug_scripts/compare_old_vs_current.py b/docs/archive/debug_scripts/compare_old_vs_current.py deleted file mode 100644 index 5645cd2..0000000 --- a/docs/archive/debug_scripts/compare_old_vs_current.py +++ /dev/null @@ -1,185 +0,0 @@ -#!/usr/bin/env python3 -"""Deep comparison: Old model (28 epoch, before fixes) vs Current model (35 epoch, after fixes).""" - -import torch -import numpy as np -import pickle - -print("=" * 80) -print("DEEP COMPARISON: OLD (28 epoch) vs CURRENT (35 epoch)") -print("=" * 80) - -# Load old model (before fixes) -old_path = "small_model_28epoch/taxonomy_model_small_best.pth" -print(f"\n📦 OLD MODEL (before fixes)") -print(f" Path: {old_path}") - -old_model = torch.load(old_path, map_location='cpu') -embs_old = old_model['embeddings'].detach().numpy() -old_epoch = old_model['epoch'] -old_loss = old_model['loss'] - -print(f" Epoch: {old_epoch}") -print(f" Loss: {old_loss:.6f}") -print(f" Shape: {embs_old.shape}") - -# Load current model (after fixes) -current_path = "taxonomy_model_small_best.pth" -print(f"\n📦 CURRENT MODEL (after fixes)") -print(f" Path: {current_path}") - -current_model = torch.load(current_path, map_location='cpu') -embs_current = current_model['embeddings'].detach().numpy() -current_epoch = current_model['epoch'] -current_loss = current_model['loss'] - -print(f" Epoch: {current_epoch}") -print(f" Loss: {current_loss:.6f}") -print(f" Shape: {embs_current.shape}") - -# Load training data to identify trained nodes -print("\n📊 Loading training info...") -with open('data/taxonomy_edges_small_transitive.pkl', 'rb') as f: - training_data = pickle.load(f) - -ancestors = set(item['ancestor_idx'] for item in training_data) -descendants = set(item['descendant_idx'] for item in training_data) -trained_indices = ancestors | descendants - -print(f" Nodes in training data: {len(trained_indices):,}") - -# Calculate norms -norms_old = np.linalg.norm(embs_old, axis=1) -norms_current = np.linalg.norm(embs_current, axis=1) - -# For OLD: Only look at nodes that exist -old_n_nodes = embs_old.shape[0] -old_trained = [i for i in trained_indices if i < old_n_nodes] -norms_old_trained = norms_old[old_trained] - -# For CURRENT: Separate trained vs untrained -current_trained = [i for i in trained_indices if i < embs_current.shape[0]] -current_untrained = [i for i in range(embs_current.shape[0]) if i not in trained_indices] -norms_current_trained = norms_current[current_trained] -norms_current_untrained = norms_current[current_untrained] if current_untrained else np.array([]) - -print("\n" + "=" * 80) -print("COMPARISON: TRAINED NODES ONLY") -print("=" * 80) - -print(f"\n📊 OLD MODEL (trained nodes only, {len(old_trained):,} nodes):") -print(f" Min norm: {norms_old_trained.min():.4f}") -print(f" Mean norm: {norms_old_trained.mean():.4f}") -print(f" Max norm: {norms_old_trained.max():.4f}") -print(f" Std dev: {norms_old_trained.std():.4f}") -print(f" >0.90: {(norms_old_trained > 0.90).sum():,} ({100*(norms_old_trained > 0.90).sum()/len(norms_old_trained):.1f}%)") -print(f" >0.95: {(norms_old_trained > 0.95).sum():,} ({100*(norms_old_trained > 0.95).sum()/len(norms_old_trained):.1f}%)") - -print(f"\n📊 CURRENT MODEL (trained nodes only, {len(current_trained):,} nodes):") -print(f" Min norm: {norms_current_trained.min():.4f}") -print(f" Mean norm: {norms_current_trained.mean():.4f}") -print(f" Max norm: {norms_current_trained.max():.4f}") -print(f" Std dev: {norms_current_trained.std():.4f}") -print(f" >0.90: {(norms_current_trained > 0.90).sum():,} ({100*(norms_current_trained > 0.90).sum()/len(norms_current_trained):.1f}%)") -print(f" >0.95: {(norms_current_trained > 0.95).sum():,} ({100*(norms_current_trained > 0.95).sum()/len(norms_current_trained):.1f}%)") - -# Variance analysis (spread) -print("\n📊 EMBEDDING SPREAD (Variance):") -print(f" OLD: {embs_old.var():.6f}") -print(f" CURRENT (all): {embs_current.var():.6f}") -print(f" CURRENT (trained only): {embs_current[current_trained].var():.6f}") - -# Percentiles -print("\n📊 NORM PERCENTILES (Trained nodes only):") -percentiles = [25, 50, 75, 90, 95, 99] -print(f"\n {'Percentile':<12} {'OLD':<10} {'CURRENT':<10} {'Δ':<10}") -print(f" {'-'*12} {'-'*10} {'-'*10} {'-'*10}") -for p in percentiles: - old_val = np.percentile(norms_old_trained, p) - current_val = np.percentile(norms_current_trained, p) - delta = current_val - old_val - print(f" {p}%{'':<10} {old_val:.4f} {current_val:.4f} {delta:+.4f}") - -# Loss comparison -print("\n📊 TRAINING LOSS:") -print(f" OLD (epoch {old_epoch}): {old_loss:.6f}") -print(f" CURRENT (epoch {current_epoch}): {current_loss:.6f}") -print(f" Δ: {current_loss - old_loss:+.6f} ({100*(current_loss - old_loss)/old_loss:+.1f}%)") - -print("\n" + "=" * 80) -print("DIAGNOSIS") -print("=" * 80) - -# Check if embeddings are more compressed -mean_diff = norms_current_trained.mean() - norms_old_trained.mean() -p90_diff = np.percentile(norms_current_trained, 90) - np.percentile(norms_old_trained, 90) -boundary_old = (norms_old_trained > 0.90).sum() / len(norms_old_trained) -boundary_current = (norms_current_trained > 0.90).sum() / len(norms_current_trained) - -issues = [] - -if mean_diff > 0.05: - issues.append(f"⚠️ Mean norm increased by {mean_diff:.3f} - embeddings pushed to boundary") - -if p90_diff > 0.05: - issues.append(f"⚠️ 90th percentile increased by {p90_diff:.3f} - more compression") - -if boundary_current > boundary_old + 0.1: - issues.append(f"⚠️ Boundary clustering increased by {100*(boundary_current - boundary_old):.1f}%") - -if norms_current_trained.std() < norms_old_trained.std() * 0.8: - issues.append(f"⚠️ Std dev decreased - less spread in radial direction") - -if embs_current[current_trained].var() < embs_old.var() * 0.8: - issues.append(f"⚠️ Total variance decreased - embeddings more clustered") - -if issues: - print("\n🔴 PROBLEMS FOUND:\n") - for issue in issues: - print(f" {issue}") -else: - print("\n✅ CURRENT model looks comparable or better") - -print("\n💡 LIKELY CAUSES:") - -# Check hyperparameters -print("\n1. **Regularization strength**") -print(" - If λ increased, embeddings get pushed to target radii") -print(" - Check: Was λ=0.1 → 0.01 → back to higher?") - -print("\n2. **Initialization**") -print(" - If init range changed, nodes start closer to boundary") -print(" - OLD init: likely [0.1, 0.95]") -print(" - CURRENT init: [0.05, 0.85]") - -print("\n3. **Number of nodes**") -print(" - OLD: {old_n_nodes:,} nodes") -print(f" - CURRENT: {embs_current.shape[0]:,} nodes") -print(" - More nodes = more crowding") - -print("\n4. **Training data**") -print(f" - OLD: {len(training_data) - 26590:,} pairs (approx)") -print(f" - CURRENT: {len(training_data):,} pairs") -print(" - More pairs might cause different dynamics") - -print("\n" + "=" * 80) -print("RECOMMENDATIONS") -print("=" * 80) - -print("\n1. **Revert to OLD hyperparameters temporarily**") -print(" - Use same λ, init range, max_norm as old model") -print(" - This isolates whether the data fix is the problem") - -print("\n2. **Check what changed between models**") -print(" - Compare train_small.py and train_hierarchical.py") -print(" - Look for λ, init range, max_norm, projection frequency") - -print("\n3. **Train longer with weaker regularization**") -print(" - Current λ might be too strong") -print(" - Try λ=0.01 for 100+ epochs") - -print("\n4. **Adjust initialization**") -print(" - Current range [0.05, 0.85] might be too conservative") -print(" - Try [0.1, 0.90] to match old model") - -print("\n" + "=" * 80) diff --git a/docs/archive/debug_scripts/diagnose_issues.py b/docs/archive/debug_scripts/diagnose_issues.py deleted file mode 100644 index ed98e1d..0000000 --- a/docs/archive/debug_scripts/diagnose_issues.py +++ /dev/null @@ -1,123 +0,0 @@ -#!/usr/bin/env python3 -"""Comprehensive diagnosis of training issues.""" - -import pickle -import pandas as pd -import numpy as np - -print("=" * 80) -print("COMPREHENSIVE DIAGNOSIS") -print("=" * 80) - -# Load training data -with open('data/taxonomy_edges_small_transitive.pkl', 'rb') as f: - training_data = pickle.load(f) - -# Load mapping -df = pd.read_csv("data/taxonomy_edges_small.mapping.tsv", sep="\t", header=None, names=["taxid", "idx"]) -df['idx'] = pd.to_numeric(df['idx'], errors='coerce') -df = df.dropna() -df['idx'] = df['idx'].astype(int) - -all_indices_in_mapping = set(df['idx']) -print(f"\nAll indices in mapping file: {len(all_indices_in_mapping):,}") -print(f" Range: {min(all_indices_in_mapping)} to {max(all_indices_in_mapping)}") - -# Get nodes from training data (what train_small.py does) -ancestor_indices = set(item['ancestor_idx'] for item in training_data) -descendant_indices = set(item['descendant_idx'] for item in training_data) -all_indices_in_training = ancestor_indices | descendant_indices - -n_nodes_from_training = max(all_indices_in_training) + 1 # What train_small.py uses -n_nodes_from_mapping = max(all_indices_in_mapping) + 1 # What it SHOULD use - -print(f"\nIndices appearing in training data: {len(all_indices_in_training):,}") -print(f" Range: {min(all_indices_in_training)} to {max(all_indices_in_training)}") -print(f" n_nodes calculated (max+1): {n_nodes_from_training:,}") - -print(f"\nExpected n_nodes from mapping: {n_nodes_from_mapping:,}") -print(f"MISMATCH: {n_nodes_from_mapping - n_nodes_from_training:,} nodes missing from model!") - -# Find missing indices -missing_indices = all_indices_in_mapping - all_indices_in_training -print(f"\nMissing {len(missing_indices):,} indices:") -print(f" They exist in mapping but never appear in training data") - -# Check what these missing indices are -only_ancestors = ancestor_indices - descendant_indices -only_descendants = descendant_indices - ancestor_indices -both = ancestor_indices & descendant_indices - -print(f"\nNode roles in training data:") -print(f" Only ancestors (never descendants): {len(only_ancestors):,}") -print(f" Only descendants (never ancestors): {len(only_descendants):,}") -print(f" Both: {len(both):,}") -print(f" Total in training: {len(all_indices_in_training):,}") -print(f" Never appear: {len(missing_indices):,}") - -# Analyze missing nodes -if len(missing_indices) < 100: - print(f"\nMissing indices: {sorted(missing_indices)[:20]}") -else: - print(f"\nFirst 20 missing indices: {sorted(missing_indices)[:20]}") - print(f"Last 20 missing indices: {sorted(missing_indices)[-20:]}") - -# Check gaps in training data indices -max_idx_training = max(all_indices_in_training) -expected_range = set(range(max_idx_training + 1)) -gaps_in_training = expected_range - all_indices_in_training - -print(f"\nGaps in training indices (0 to {max_idx_training}):") -print(f" {len(gaps_in_training):,} indices missing from continuous range") -if len(gaps_in_training) < 50: - print(f" Gaps: {sorted(gaps_in_training)}") - -print("\n" + "=" * 80) -print("ROOT CAUSE ANALYSIS") -print("=" * 80) - -print("\n1. BUG IN train_small.py (line 346-347):") -print(" Current: n_nodes = max(training_data indices) + 1") -print(" Problem: Excludes nodes not in training pairs") -print(" Fix: n_nodes = max(mapping indices) + 1") - -print("\n2. BOUNDARY COMPRESSION:") -print(" 90% of embeddings at norm > 0.9997") -print(" Problem: No space for hierarchical structure") -print(" Causes:") -print(" - Radial regularizer pushing to specific radii") -print(" - Ball projection clamping at 1.0") -print(" - Deep pairs (74%) all pushed near boundary") - -print("\n3. DATA IMBALANCE:") -print(" Parent-child: 6%, Deep pairs (>5): 74%") -print(" Problem: Model overfits to deep relationships") -print(" Effect: Local structure (siblings, parents) lost") - -print("\n" + "=" * 80) -print("RECOMMENDED FIXES") -print("=" * 80) - -print("\n1. CRITICAL - Fix n_nodes calculation:") -print(" Replace lines 346-347 in train_small.py with:") -print(" ```") -print(" # Load mapping to get true n_nodes") -print(" mapping_df = pd.read_csv('data/taxonomy_edges_small.mapping.tsv',") -print(" sep='\\t', header=None, names=['taxid', 'idx'])") -print(" mapping_df['idx'] = pd.to_numeric(mapping_df['idx'], errors='coerce')") -print(" mapping_df = mapping_df.dropna()") -print(" n_nodes = int(mapping_df['idx'].max()) + 1") -print(" ```") - -print("\n2. URGENT - Relax boundary constraint:") -print(" Change project_to_ball max radius from 0.999.. to 0.98") -print(" Leave room for hierarchical spread") - -print("\n3. IMPORTANT - Balance training data:") -print(" Option A: Oversample parent-child pairs (weight them 10x)") -print(" Option B: Cap depth to 5-10 levels") -print(" Option C: Curriculum learning (start with depth 1-2, gradually add deeper)") - -print("\n4. OPTIONAL - Reduce radial regularization:") -print(" Try lambda_reg = 0.01 instead of 0.1") -print(" Or disable completely for first epochs") diff --git a/docs/archive/debug_scripts/find_what_broke.py b/docs/archive/debug_scripts/find_what_broke.py deleted file mode 100644 index b19f409..0000000 --- a/docs/archive/debug_scripts/find_what_broke.py +++ /dev/null @@ -1,80 +0,0 @@ -#!/usr/bin/env python3 -"""Find what hyperparameter changes broke the model.""" - -print("=" * 80) -print("WHAT CHANGED BETWEEN OLD AND CURRENT") -print("=" * 80) - -print("\n📊 DATA:") -print(" OLD: 975,896 training pairs") -print(" CURRENT: 1,002,486 training pairs (+26,590)") -print(" ✅ This is GOOD - more complete coverage") - -print("\n📊 MY CHANGES (that might have broken things):") - -print("\n1. **Regularization strength (λ)**") -print(" OLD: λ = 0.1 (default before my changes)") -print(" CURRENT: λ = 0.01 (I reduced it 10x)") -print(" 🔴 TOO WEAK - not enough structure enforcement") - -print("\n2. **Initialization range**") -print(" OLD: target_radius = 0.1 + depth/max * 0.85") -print(" → Range: [0.1, 0.95]") -print(" CURRENT: target_radius = 0.05 + depth/max * 0.80") -print(" → Range: [0.05, 0.85]") -print(" 🔴 MORE CONSERVATIVE - less room to learn") - -print("\n3. **Hard projection boundary**") -print(" OLD: max_norm = 0.99999 (essentially 1.0)") -print(" CURRENT: max_norm = 0.98") -print(" 🔴 TOO TIGHT - embeddings get compressed") - -print("\n4. **Number of nodes**") -print(" OLD: 92,290 nodes (incomplete)") -print(" CURRENT: 111,103 nodes (complete)") -print(" ⚠️ More nodes = more crowding, but this is necessary") - -print("\n" + "=" * 80) -print("THE PROBLEM") -print("=" * 80) - -print(""" -I made THREE hyperparameter changes that were TOO conservative: - -1. λ: 0.1 → 0.01 (10x weaker) - → Not enough regularization for 111K nodes - -2. Init: [0.1, 0.95] → [0.05, 0.85] - → Starts too far from boundary, harder to learn hierarchy - -3. max_norm: 0.99999 → 0.98 - → Artificial compression near boundary - -These changes were meant to "fix" boundary compression, but they were -based on the OLD incomplete data. With complete data, the old hyperparams -were actually fine! - -The REAL issue was just the missing 18,813 nodes, which we now fixed. -But I over-corrected the hyperparameters. -""") - -print("\n" + "=" * 80) -print("THE FIX") -print("=" * 80) - -print(""" -KEEP: Complete data coverage (98.3% = excellent) -REVERT: All three hyperparameter changes - -Specifically: -1. λ = 0.1 (not 0.01) -2. Init range: [0.1, 0.95] (not [0.05, 0.85]) -3. max_norm = 0.999 (not 0.98) - -This gives us: -✅ Complete data coverage (the real permanent fix) -✅ Original hyperparameters that worked well -✅ Should match or beat OLD model performance -""") - -print("\n" + "=" * 80) diff --git a/docs/archive/debug_scripts/inspect_checkpoint.py b/docs/archive/debug_scripts/inspect_checkpoint.py deleted file mode 100644 index 0363af1..0000000 --- a/docs/archive/debug_scripts/inspect_checkpoint.py +++ /dev/null @@ -1,39 +0,0 @@ -#!/usr/bin/env python3 -"""Inspect checkpoint structure.""" - -import torch - -old_ckpt = "taxonomy_model_small_epoch57.pth" -new_ckpt = "taxonomy_model_small_epoch36.pth" - -print("=" * 80) -print("OLD CHECKPOINT STRUCTURE") -print("=" * 80) - -model_old = torch.load(old_ckpt, map_location='cpu') -print(f"\nType: {type(model_old)}") -if isinstance(model_old, dict): - print(f"Keys: {list(model_old.keys())}") - for key in model_old.keys(): - val = model_old[key] - print(f"\n '{key}': {type(val)}") - if isinstance(val, dict): - print(f" Sub-keys: {list(val.keys())[:10]}") - elif hasattr(val, 'shape'): - print(f" Shape: {val.shape}") - -print("\n" + "=" * 80) -print("NEW CHECKPOINT STRUCTURE") -print("=" * 80) - -model_new = torch.load(new_ckpt, map_location='cpu') -print(f"\nType: {type(model_new)}") -if isinstance(model_new, dict): - print(f"Keys: {list(model_new.keys())}") - for key in model_new.keys(): - val = model_new[key] - print(f"\n '{key}': {type(val)}") - if isinstance(val, dict): - print(f" Sub-keys: {list(val.keys())[:10]}") - elif hasattr(val, 'shape'): - print(f" Shape: {val.shape}") diff --git a/docs/archive/debug_scripts/test_depth_coverage.py b/docs/archive/debug_scripts/test_depth_coverage.py deleted file mode 100644 index 3f2e445..0000000 --- a/docs/archive/debug_scripts/test_depth_coverage.py +++ /dev/null @@ -1,76 +0,0 @@ -#!/usr/bin/env python3 -"""Quick test to verify depth coverage fix.""" - -import pandas as pd -import pickle - -print("Testing depth coverage fix...\n") - -# Load training data -with open('data/taxonomy_edges_small_transitive.pkl', 'rb') as f: - training_data = pickle.load(f) - -# Load mapping -mapping_df = pd.read_csv("data/taxonomy_edges_small.mapping.tsv", - sep="\t", header=None, names=["taxid", "idx"]) -mapping_df['idx'] = pd.to_numeric(mapping_df['idx'], errors='coerce') -mapping_df['taxid'] = pd.to_numeric(mapping_df['taxid'], errors='coerce') -mapping_df = mapping_df.dropna() -mapping_df['idx'] = mapping_df['idx'].astype(int) -mapping_df['taxid'] = mapping_df['taxid'].astype(int) -n_nodes = int(mapping_df['idx'].max()) + 1 - -max_depth = max(item['descendant_depth'] for item in training_data) - -print(f"Total nodes in dataset: {n_nodes:,}") -print(f"Max depth: {max_depth}\n") - -# Build depth mapping from training data (old way) -idx_to_depth_old = {} -for item in training_data: - idx_to_depth_old[item['descendant_idx']] = item['descendant_depth'] - if item['ancestor_idx'] not in idx_to_depth_old: - idx_to_depth_old[item['ancestor_idx']] = item['ancestor_depth'] - -print(f"OLD approach: {len(idx_to_depth_old):,} nodes with depths") - -# Build TaxID -> depth mapping -taxid_to_depth = {} -for item in training_data: - taxid_to_depth[item['ancestor_taxid']] = item['ancestor_depth'] - taxid_to_depth[item['descendant_taxid']] = item['descendant_depth'] - -print(f"TaxID -> depth mapping: {len(taxid_to_depth):,} TaxIDs") - -# NEW approach: fill in missing nodes -idx_to_depth_new = dict(idx_to_depth_old) -missing_count = 0 -assigned_from_taxid = 0 -assigned_default = 0 - -for idx in range(n_nodes): - if idx not in idx_to_depth_new: - missing_count += 1 - # Find TaxID for this index - taxid = mapping_df[mapping_df['idx'] == idx]['taxid'].values[0] - if taxid in taxid_to_depth: - idx_to_depth_new[idx] = taxid_to_depth[taxid] - assigned_from_taxid += 1 - else: - # Node not in taxonomy at all - likely a leaf, assign max depth - idx_to_depth_new[idx] = max_depth - assigned_default += 1 - -print(f"\nNEW approach: {len(idx_to_depth_new):,} nodes with depths") -print(f" Missing from training: {missing_count:,}") -print(f" Assigned via TaxID lookup: {assigned_from_taxid:,}") -print(f" Assigned default (max_depth): {assigned_default:,}") - -if len(idx_to_depth_new) == n_nodes: - print(f"\n✅ SUCCESS: All {n_nodes:,} nodes have depth information!") -else: - print(f"\n❌ PROBLEM: Only {len(idx_to_depth_new):,} / {n_nodes:,} nodes have depths") - -print(f"\nExpected regularization message:") -print(f" OLD: 'Regularizing {len(idx_to_depth_old):,} nodes'") -print(f" NEW: 'Regularizing {len(idx_to_depth_new):,} nodes'") diff --git a/docs/archive/debug_scripts/verify_ball_safety.py b/docs/archive/debug_scripts/verify_ball_safety.py deleted file mode 100644 index ed3e4c3..0000000 --- a/docs/archive/debug_scripts/verify_ball_safety.py +++ /dev/null @@ -1,188 +0,0 @@ -#!/usr/bin/env python3 -"""Verify that all ball constraint protections are still active.""" - -import ast -import re - -print("=" * 80) -print("BALL CONSTRAINT SAFETY VERIFICATION") -print("=" * 80) - -# Read train_small.py -with open('train_small.py', 'r') as f: - train_small_code = f.read() - -# Read train_hierarchical.py -with open('train_hierarchical.py', 'r') as f: - train_hierarchical_code = f.read() - -print("\n1. CHECKING GRADIENT CLIPPING...") -if 'clip_grad_norm_' in train_small_code and 'max_norm=1.0' in train_small_code: - print(" ✅ Gradient clipping ACTIVE (max_norm=1.0)") - print(" Location: train_small.py line 203") -else: - print(" ❌ WARNING: Gradient clipping NOT FOUND") - -print("\n2. CHECKING PER-BATCH PROJECTION...") -per_batch_match = re.search(r'model\.project_to_ball\(updated_indices\)', train_small_code) -if per_batch_match: - print(" ✅ Per-batch projection ACTIVE") - print(" Location: train_small.py line 210") -else: - print(" ❌ WARNING: Per-batch projection NOT FOUND") - -print("\n3. CHECKING PERIODIC PROJECTION...") -periodic_match = re.search(r'if n_batches % 500.*?model\.project_to_ball\(indices=None\)', - train_small_code, re.DOTALL) -if periodic_match: - print(" ✅ Periodic projection ACTIVE (every 500 batches)") - print(" Location: train_small.py lines 212-214") -else: - print(" ❌ WARNING: Periodic projection NOT FOUND") - -print("\n4. CHECKING EPOCH-END PROJECTION...") -epoch_match = re.search(r'# Final projection at epoch end.*?model\.project_to_ball\(indices=None\)', - train_small_code, re.DOTALL) -if epoch_match: - print(" ✅ Epoch-end projection ACTIVE") - print(" Location: train_small.py line 227") -else: - print(" ❌ WARNING: Epoch-end projection NOT FOUND") - -print("\n5. CHECKING PROJECTION IMPLEMENTATION...") -# Check max_norm parameter -if 'def project_to_ball(self, indices=None, max_norm=0.98)' in train_hierarchical_code: - print(" ✅ Hard constraint: max_norm = 0.98") - print(" Location: train_hierarchical.py line 104") - - # Verify hard projection logic - if 'needs_projection = norms >= max_norm' in train_hierarchical_code: - print(" ✅ Hard projection logic: only scales violators") - else: - print(" ⚠️ Projection logic may have changed") -else: - print(" ❌ WARNING: max_norm parameter not found or changed") - -print("\n6. CHECKING REGULARIZATION STRENGTH...") -if '--lambda-reg' in train_small_code: - # Extract default value - lambda_match = re.search(r"'--lambda-reg'.*?default=([\d.]+)", train_small_code) - if lambda_match: - lambda_val = float(lambda_match.group(1)) - print(f" ✅ Regularization strength: λ = {lambda_val}") - if lambda_val == 0.01: - print(" ℹ️ Reduced from 0.1 → 0.01 (10x weaker, but safe with projections)") - elif lambda_val == 0.1: - print(" ℹ️ Original strong value (0.1)") - else: - print(" ⚠️ Could not parse lambda value") - -print("\n" + "=" * 80) -print("COMPARISON TO HISTORICAL VERSIONS") -print("=" * 80) - -versions = { - 'v1 (Broken)': { - 'lambda': 0.01, - 'grad_clip': False, - 'per_batch_proj': False, - 'periodic_proj': False, - 'epoch_proj': False, - 'max_norm': None, - 'result': '54% escaped' - }, - 'v2 (Better)': { - 'lambda': 0.1, - 'grad_clip': True, - 'per_batch_proj': False, - 'periodic_proj': False, - 'epoch_proj': False, - 'max_norm': 0.99999, - 'result': '2.2% escaped' - }, - 'v3 (Fixed)': { - 'lambda': 0.1, - 'grad_clip': True, - 'per_batch_proj': True, - 'periodic_proj': True, - 'epoch_proj': True, - 'max_norm': 0.99999, - 'result': '0% escaped ✅' - }, - 'Current (Our fixes)': { - 'lambda': 0.01, - 'grad_clip': True, - 'per_batch_proj': True, - 'periodic_proj': True, - 'epoch_proj': True, - 'max_norm': 0.98, - 'result': 'Expected: 0% ✅' - } -} - -print("\n| Feature | v1 | v2 | v3 | Current |") -print("|---------|----|----|----|---------| ") -print(f"| Lambda | {versions['v1 (Broken)']['lambda']} | {versions['v2 (Better)']['lambda']} | {versions['v3 (Fixed)']['lambda']} | {versions['Current (Our fixes)']['lambda']} |") -print(f"| Grad clip | {'✅' if versions['v1 (Broken)']['grad_clip'] else '❌'} | {'✅' if versions['v2 (Better)']['grad_clip'] else '❌'} | {'✅' if versions['v3 (Fixed)']['grad_clip'] else '❌'} | {'✅' if versions['Current (Our fixes)']['grad_clip'] else '❌'} |") -print(f"| Per-batch proj | {'✅' if versions['v1 (Broken)']['per_batch_proj'] else '❌'} | {'✅' if versions['v2 (Better)']['per_batch_proj'] else '❌'} | {'✅' if versions['v3 (Fixed)']['per_batch_proj'] else '❌'} | {'✅' if versions['Current (Our fixes)']['per_batch_proj'] else '❌'} |") -print(f"| Periodic proj | {'✅' if versions['v1 (Broken)']['periodic_proj'] else '❌'} | {'✅' if versions['v2 (Better)']['periodic_proj'] else '❌'} | {'✅' if versions['v3 (Fixed)']['periodic_proj'] else '❌'} | {'✅' if versions['Current (Our fixes)']['periodic_proj'] else '❌'} |") -print(f"| Epoch proj | {'✅' if versions['v1 (Broken)']['epoch_proj'] else '❌'} | {'✅' if versions['v2 (Better)']['epoch_proj'] else '❌'} | {'✅' if versions['v3 (Fixed)']['epoch_proj'] else '❌'} | {'✅' if versions['Current (Our fixes)']['epoch_proj'] else '❌'} |") -print(f"| max_norm | {versions['v1 (Broken)']['max_norm'] or 'None'} | {versions['v2 (Better)']['max_norm']} | {versions['v3 (Fixed)']['max_norm']} | **{versions['Current (Our fixes)']['max_norm']}** |") -print(f"| **Result** | {versions['v1 (Broken)']['result']} | {versions['v2 (Better)']['result']} | {versions['v3 (Fixed)']['result']} | **{versions['Current (Our fixes)']['result']}** |") - -print("\n" + "=" * 80) -print("SAFETY ASSESSMENT") -print("=" * 80) - -protections = [ - ("Gradient clipping", True), - ("Per-batch projection", True), - ("Periodic projection (500 batches)", True), - ("Epoch-end projection", True), - ("Hard max_norm constraint", True) -] - -active_count = sum(1 for _, status in protections if status) -print(f"\nActive protections: {active_count}/{len(protections)}") -for name, status in protections: - symbol = "✅" if status else "❌" - print(f" {symbol} {name}") - -print("\n" + "=" * 80) -print("KEY INSIGHT") -print("=" * 80) -print(""" -The ball escape bug (v1) was caused by LACK of hard constraints. -We had only soft regularization (λ=0.01) with no projection. - -The fix (v3) added 3-layer HARD PROJECTION strategy. -This GUARANTEES embeddings stay inside, regardless of regularization. - -Our changes: - - Reduced λ: 0.1 → 0.01 (weaker soft guidance) - - Reduced max_norm: 0.99999 → 0.98 (STRICTER hard limit!) - -Result: SAFER than v3 - - More room from boundary (2% vs 0.001%) - - Still impossible to escape (hard projection enforces it) - - Better spread (less compression at boundary) -""") - -print("\n" + "=" * 80) -print("VERDICT") -print("=" * 80) - -if active_count == len(protections): - print("\n✅ ✅ ✅ ALL SAFETY MECHANISMS ACTIVE ✅ ✅ ✅") - print("\nBall escape is MATHEMATICALLY IMPOSSIBLE with:") - print(" 1. Hard projection enforcing max_norm=0.98") - print(" 2. Three projection layers (batch, periodic, epoch)") - print(" 3. Gradient clipping preventing explosive updates") - print("\nRegularization (λ=0.01) only affects HOW SMOOTHLY we learn,") - print("not WHETHER we stay inside the ball.") - print("\n🟢 SAFE TO PROCEED - Risk level: NEGLIGIBLE") -else: - print(f"\n⚠️ WARNING: Only {active_count}/{len(protections)} protections active!") - print("Review code before training.") - -print("\n" + "=" * 80) diff --git a/docs/archive/debug_scripts/verify_fixes.py b/docs/archive/debug_scripts/verify_fixes.py deleted file mode 100644 index 6c19dfb..0000000 --- a/docs/archive/debug_scripts/verify_fixes.py +++ /dev/null @@ -1,136 +0,0 @@ -#!/usr/bin/env python3 -"""Verify that all fixes are properly applied.""" - -import pickle -import pandas as pd - -print("=" * 80) -print("VERIFICATION: All Fixes Applied Correctly") -print("=" * 80) - -# Test 1: n_nodes calculation -print("\n1. Testing n_nodes calculation...") -with open('data/taxonomy_edges_small_transitive.pkl', 'rb') as f: - training_data = pickle.load(f) - -# OLD WAY (buggy) -old_n_nodes = max(max(item['ancestor_idx'], item['descendant_idx']) - for item in training_data) + 1 - -# NEW WAY (correct) -mapping_df = pd.read_csv("data/taxonomy_edges_small.mapping.tsv", - sep="\t", header=None, names=["taxid", "idx"]) -mapping_df['idx'] = pd.to_numeric(mapping_df['idx'], errors='coerce') -mapping_df = mapping_df.dropna() -new_n_nodes = int(mapping_df['idx'].max()) + 1 - -print(f" OLD calculation (buggy): {old_n_nodes:,} nodes") -print(f" NEW calculation (fixed): {new_n_nodes:,} nodes") -print(f" Difference: {new_n_nodes - old_n_nodes:,} nodes recovered") - -if new_n_nodes > old_n_nodes: - print(f" ✅ FIX VERIFIED: Will create {new_n_nodes:,} embeddings (covering all nodes)") -else: - print(f" ❌ PROBLEM: n_nodes unchanged") - -# Test 2: Check initialization range -print("\n2. Testing initialization range...") -max_depth = 37 - -# OLD initialization -old_min_radius = 0.1 + (0 / max_depth) * 0.85 -old_max_radius = 0.1 + (max_depth / max_depth) * 0.85 - -# NEW initialization -new_min_radius = 0.05 + (0 / max_depth) * 0.80 -new_max_radius = 0.05 + (max_depth / max_depth) * 0.80 - -print(f" OLD range: [{old_min_radius:.2f}, {old_max_radius:.2f}]") -print(f" NEW range: [{new_min_radius:.2f}, {new_max_radius:.2f}]") -print(f" Buffer from boundary (1.0): {1.0 - new_max_radius:.2f}") - -if new_max_radius < old_max_radius: - print(f" ✅ FIX VERIFIED: {100*(1.0-new_max_radius):.0f}% buffer from boundary") -else: - print(f" ❌ PROBLEM: Still too close to boundary") - -# Test 3: Check projection constraint -print("\n3. Testing projection constraint...") -old_max_norm = 1.0 - 1e-5 -new_max_norm = 0.98 - -print(f" OLD max_norm: {old_max_norm:.6f}") -print(f" NEW max_norm: {new_max_norm:.6f}") -print(f" Buffer gained: {old_max_norm - new_max_norm:.6f}") - -if new_max_norm < old_max_norm: - print(f" ✅ FIX VERIFIED: {100*(1.0-new_max_norm):.0f}% buffer from boundary") -else: - print(f" ⚠️ Check: max_norm should be 0.98 in code") - -# Test 4: Regularization strength -print("\n4. Testing regularization strength...") -old_lambda = 0.1 -new_lambda = 0.01 - -print(f" OLD lambda_reg: {old_lambda}") -print(f" NEW lambda_reg: {new_lambda}") -print(f" Reduction: {old_lambda / new_lambda:.0f}x weaker") - -if new_lambda < old_lambda: - print(f" ✅ FIX VERIFIED: {old_lambda/new_lambda:.0f}x weaker regularization") -else: - print(f" ❌ PROBLEM: lambda_reg should be reduced") - -# Test 5: Data coverage -print("\n5. Testing data coverage...") -ancestor_indices = set(item['ancestor_idx'] for item in training_data) -descendant_indices = set(item['descendant_idx'] for item in training_data) -all_in_training = ancestor_indices | descendant_indices -all_in_mapping = set(mapping_df['idx'].unique()) - -covered_old = len(all_in_training) -covered_new = len(all_in_mapping) -coverage_pct = 100 * covered_old / covered_new - -print(f" Nodes in training data: {covered_old:,}") -print(f" Nodes in mapping: {covered_new:,}") -print(f" OLD coverage: {coverage_pct:.1f}%") -print(f" NEW coverage: 100.0% (all nodes get embeddings)") - -missing = all_in_mapping - all_in_training -print(f" Nodes that will now be initialized (were missing): {len(missing):,}") - -if len(missing) > 0: - print(f" ✅ FIX VERIFIED: {len(missing):,} additional nodes will get proper embeddings") -else: - print(f" ❓ No missing nodes found") - -# Summary -print("\n" + "=" * 80) -print("SUMMARY") -print("=" * 80) - -checks = [] -checks.append(("n_nodes calculation", new_n_nodes > old_n_nodes)) -checks.append(("Initialization range", new_max_radius < old_max_radius)) -checks.append(("Projection constraint", new_max_norm < old_max_norm)) -checks.append(("Regularization strength", new_lambda < old_lambda)) -checks.append(("Data coverage", len(missing) > 0)) - -passed = sum(1 for _, status in checks if status) -total = len(checks) - -print(f"\nChecks passed: {passed}/{total}") -for name, status in checks: - symbol = "✅" if status else "❌" - print(f" {symbol} {name}") - -if passed == total: - print("\n🎉 ALL FIXES VERIFIED - Ready to retrain!") - print("\nRecommended command:") - print(" uv run python train_small.py --epochs 10000 --early-stopping 0") -else: - print(f"\n⚠️ {total - passed} checks failed - review fixes") - -print("\n" + "=" * 80) diff --git a/docs/archive/debug_scripts/visualize_trained_only.py b/docs/archive/debug_scripts/visualize_trained_only.py deleted file mode 100644 index 59d0b70..0000000 --- a/docs/archive/debug_scripts/visualize_trained_only.py +++ /dev/null @@ -1,187 +0,0 @@ -#!/usr/bin/env python3 -"""Visualize only the nodes that were actually trained (appeared in training pairs).""" - -import torch -import numpy as np -import pandas as pd -import pickle -import umap -import matplotlib.pyplot as plt -from collections import defaultdict -import sys - -print("=" * 80) -print("VISUALIZATION: TRAINED NODES ONLY") -print("=" * 80) - -# Load checkpoint -checkpoint_file = sys.argv[1] if len(sys.argv) > 1 else "taxonomy_model_small_best.pth" -print(f"\nLoading embeddings from {checkpoint_file}...") - -model = torch.load(checkpoint_file, map_location='cpu') -if isinstance(model, dict) and 'embeddings' in model: - embeddings = model['embeddings'].detach().numpy() -else: - print("Error: Unexpected checkpoint format") - sys.exit(1) - -print(f" ✓ Total embeddings: {embeddings.shape}") - -# Load training data to identify trained nodes -print("\nLoading training data to identify trained nodes...") -with open('data/taxonomy_edges_small_transitive.pkl', 'rb') as f: - training_data = pickle.load(f) - -ancestors = set(item['ancestor_idx'] for item in training_data) -descendants = set(item['descendant_idx'] for item in training_data) -trained_indices = sorted(ancestors | descendants) - -print(f" ✓ Nodes in training: {len(trained_indices):,}") -print(f" ✗ Nodes NOT in training: {embeddings.shape[0] - len(trained_indices):,}") - -# Filter to trained nodes only -embeddings_trained = embeddings[trained_indices] -norms = np.linalg.norm(embeddings_trained, axis=1) - -print(f"\n📊 Trained nodes statistics:") -print(f" Shape: {embeddings_trained.shape}") -print(f" Norm range: [{norms.min():.3f}, {norms.max():.3f}]") -print(f" Norm mean: {norms.mean():.3f}") -print(f" At boundary (>0.95): {(norms > 0.95).sum()} ({100*(norms > 0.95).sum()/len(norms):.1f}%)") - -# Load mapping -print("\nLoading taxonomy mapping...") -mapping_df = pd.read_csv("data/taxonomy_edges_small.mapping.tsv", - sep="\t", header=None, names=["taxid", "idx"]) -mapping_df['idx'] = pd.to_numeric(mapping_df['idx'], errors='coerce') -mapping_df['taxid'] = pd.to_numeric(mapping_df['taxid'], errors='coerce') -mapping_df = mapping_df.dropna() -mapping_df['idx'] = mapping_df['idx'].astype(int) -mapping_df['taxid'] = mapping_df['taxid'].astype(int) - -# Load taxonomy tree -print("Loading taxonomy tree...") -idx_to_taxid = dict(zip(mapping_df['idx'], mapping_df['taxid'])) -valid_taxids = set(idx_to_taxid.values()) - -# Load nodes (parent relationships) -nodes = {} -with open("data/nodes.dmp", "r") as f: - for line in f: - parts = [p.strip() for p in line.split("|")] - if len(parts) >= 2: - taxid = int(parts[0]) - parent = int(parts[1]) - if taxid in valid_taxids: - nodes[taxid] = parent - -print(f" ✓ Loaded {len(nodes):,} taxonomy nodes") - -# Identify groups -print("\nFinding taxonomic groups...") - -def find_group_descendants(root_taxid, nodes, idx_to_taxid, trained_indices): - """Find all descendants of a taxonomic group in trained indices.""" - group_indices = [] - - def find_descendants(taxid): - # Check if this taxid is in our trained nodes - for idx in trained_indices: - if idx_to_taxid.get(idx) == taxid: - group_indices.append(idx) - - # Find children - for child, parent in nodes.items(): - if parent == taxid: - find_descendants(child) - - find_descendants(root_taxid) - return set(group_indices) - -groups = { - 'Mammals': 40674, - 'Birds': 8782, - 'Insects': 50557, - 'Bacteria': 2, - 'Fungi': 4751, - 'Plants': 33090 -} - -group_indices = {} -for name, root_taxid in groups.items(): - indices = find_group_descendants(root_taxid, nodes, idx_to_taxid, trained_indices) - group_indices[name] = indices - print(f" ✓ {name}: {len(indices):,} organisms") - -# Sample if needed -sample_size = 25000 -if len(trained_indices) > sample_size: - print(f"\nSampling {sample_size:,} from {len(trained_indices):,} trained nodes...") - sample_idx = np.random.choice(len(trained_indices), sample_size, replace=False) - embeddings_plot = embeddings_trained[sample_idx] - sampled_indices = [trained_indices[i] for i in sample_idx] -else: - embeddings_plot = embeddings_trained - sampled_indices = trained_indices - sample_idx = np.arange(len(trained_indices)) - -print(f" ✓ Using {len(embeddings_plot):,} points") - -# Assign colors -colors = [] -color_map = { - 'Mammals': '#e74c3c', - 'Birds': '#f1c40f', - 'Insects': '#2ecc71', - 'Bacteria': '#9b59b6', - 'Fungi': '#e67e22', - 'Plants': '#1abc9c' -} - -for idx in sampled_indices: - assigned = False - for group_name, group_set in group_indices.items(): - if idx in group_set: - colors.append(color_map[group_name]) - assigned = True - break - if not assigned: - colors.append('#bdc3c7') - -# UMAP -print(f"\nRunning UMAP on {len(embeddings_plot):,} points...") -reducer = umap.UMAP( - n_neighbors=15, - min_dist=0.1, - metric='euclidean', - random_state=42 -) -embedding_2d = reducer.fit_transform(embeddings_plot) - -# Plot -print("Creating visualization...") -plt.figure(figsize=(16, 12)) -plt.scatter(embedding_2d[:, 0], embedding_2d[:, 1], - c=colors, s=1, alpha=0.6, rasterized=True) - -# Legend -from matplotlib.patches import Patch -legend_elements = [Patch(facecolor=color_map[name], label=f"{name} ({len(group_indices[name]):,})") - for name in groups.keys() if len(group_indices[name]) > 0] -legend_elements.append(Patch(facecolor='#bdc3c7', label='Other')) -plt.legend(handles=legend_elements, loc='upper right', fontsize=10) - -plt.title(f'Poincaré Embeddings - TRAINED NODES ONLY ({len(trained_indices):,} nodes)\n' - f'Excluding {embeddings.shape[0] - len(trained_indices):,} untrained nodes', - fontsize=16, pad=20) -plt.xlabel('UMAP 1', fontsize=12) -plt.ylabel('UMAP 2', fontsize=12) -plt.grid(True, alpha=0.3) -plt.tight_layout() - -output_file = 'taxonomy_embeddings_trained_only.png' -plt.savefig(output_file, dpi=150, bbox_inches='tight') -print(f"\n✅ Saved: {output_file}") -print(f"\n Showing: {len(trained_indices):,} trained nodes") -print(f" Hidden: {embeddings.shape[0] - len(trained_indices):,} untrained nodes") -print("\n" + "=" * 80) diff --git a/docs/theory.md b/docs/theory.md new file mode 100644 index 0000000..1b0e1ff --- /dev/null +++ b/docs/theory.md @@ -0,0 +1,250 @@ +# Theory: Poincaré Embeddings for Hierarchies + +Understanding the mathematics and intuition behind hyperbolic embeddings. + +## The Problem: Representing Hierarchies + +Traditional embeddings (Word2Vec, GloVe) use **Euclidean space** (flat space). But hierarchical data like taxonomies have **exponential growth**: + +``` +Root: 1 node +Level 1: 10 nodes +Level 2: 100 nodes +Level 3: 1,000 nodes +Level 4: 10,000 nodes +``` + +In flat Euclidean space, you need **exponentially growing dimensions** to represent this without distortion. + +## The Solution: Hyperbolic Space + +**Poincaré embeddings** use **hyperbolic geometry** where: + +- Distance grows **exponentially** as you move from center to boundary +- Perfect for hierarchies: root near center, leaves near boundary +- Can represent exponential growth in **constant dimensions** + +### The Poincaré Ball Model + +The n-dimensional Poincaré ball is: + +``` +B^n = {x ∈ R^n : ||x|| < 1} +``` + +All points inside the unit ball. The boundary (||x|| = 1) is "at infinity". + +### Poincaré Distance + +The hyperbolic distance between points u and v is: + +``` +d(u,v) = arcosh(1 + 2||u-v||²/((1-||u||²)(1-||v||²))) +``` + +Key properties: + +- Distances grow exponentially near boundary +- Center (origin) is "special" - represents root +- Angular distance represents similarity + +## Why It Works for Taxonomies + +### 1. Natural Hierarchy Encoding + +In the Poincaré ball: + +- **Root organisms** (cellular life) → center (||x|| ≈ 0.1) +- **Intermediate levels** (kingdoms, phyla) → middle (||x|| ≈ 0.5) +- **Leaf organisms** (species) → boundary (||x|| ≈ 0.95) + +Depth in taxonomy directly corresponds to norm (distance from origin). + +### 2. Exponential Capacity + +Near the boundary, even small angular differences create large hyperbolic distances. This allows: + +- Millions of species to fit in 10-20 dimensions +- Each level has exponentially more "room" than the previous + +### 3. Hierarchy Preservation + +The loss function encourages: + +- Ancestors closer to descendants than to random nodes +- Deeper pairs (great-grandparent → descendant) to have larger distances +- Siblings (same depth) to have similar norms but different angles + +## Training Methodology + +### Loss Function + +We use a **ranking loss with margin**: + +``` +L = max(0, d(ancestor, descendant) - d(ancestor, negative) + margin) +``` + +This encourages: + +- `d(ancestor, descendant)` < `d(ancestor, negative) + margin` +- Ancestors should be closer to their descendants than to random nodes + +### Radial Regularization + +We add a penalty to encourage depth → radius mapping: + +``` +R = λ Σ (||embedding_i|| - target_radius_i)² +``` + +Where `target_radius = 0.1 + (depth/max_depth) × 0.85` + +This ensures: + +- Root nodes stay near center +- Leaf nodes move toward boundary +- Smooth gradient across depths + +### Hard Negative Sampling + +Instead of random negatives, we sample **cousins** (nodes at same depth): + +- More informative: teaches model to distinguish siblings +- Depth-stratified: ensures coverage across all levels +- Harder: creates better separation + +## Optimization Challenges + +### Ball Constraint + +All embeddings must satisfy ||x|| < 1. We enforce this with: + +1. **Gradient clipping**: Prevents large jumps +2. **Selective projection**: After each batch, project updated embeddings +3. **Full projection**: At epoch end, ensure all embeddings inside ball + +### Riemannian Optimization + +Standard SGD assumes Euclidean space. For hyperbolic space, we could use: + +- **Riemannian SGD**: Updates along geodesics +- **Exponential map**: Maps tangent vectors to manifold + +Currently, we use standard SGD with careful projection - simpler and works well. + +## Theoretical Guarantees + +### Embedding Capacity + +In d-dimensional Poincaré ball, the number of points with pairwise distance ≥ δ grows as: + +``` +N ≈ exp(d × δ) +``` + +For taxonomies: + +- 10D can embed ~10^6 nodes with good separation +- 20D can embed ~10^12 nodes + +### Distortion Bounds + +For trees with branching factor b and depth h: + +- **Euclidean**: needs O(b^h) dimensions +- **Hyperbolic**: needs O(h) dimensions + +Exponential improvement! + +## Comparison to Other Approaches + +### vs. Euclidean Embeddings + +- ❌ Euclidean: Poor for hierarchies, needs many dimensions +- ✅ Hyperbolic: Natural fit, constant dimensions + +### vs. Graph Neural Networks + +- ❌ GNNs: Complex, slow, need many layers for deep trees +- ✅ Hyperbolic: Direct embedding, fast training + +### vs. Order Embeddings + +- ❌ Order: Can represent partial orders, but wastes dimensions +- ✅ Hyperbolic: Efficient, captures similarity and order + +## Practical Considerations + +### Choosing Dimensionality + +Rule of thumb: + +- **Small taxonomies** (<10K nodes): 5-10D sufficient +- **Medium** (10K-1M): 10-20D +- **Large** (>1M): 20-50D + +Higher dimensions allow more nuance but train slower. + +### Choosing Regularization + +λ controls radial constraint: + +- **Too small** (λ < 0.01): Nodes may ignore depth structure +- **Just right** (λ = 0.1): Enforces depth → radius while allowing flexibility +- **Too large** (λ > 0.5): Over-constrained, poor performance + +### Numerical Stability + +Near the boundary (||x|| ≈ 1), distances can explode. We handle this with: + +- Clamping norms: `||x|| < 0.999` (not exactly 1) +- Small epsilon in formulas: Prevents division by zero +- Gradient clipping: Avoids NaN/Inf + +## Mathematical Foundations + +### Hyperbolic Geometry + +The Poincaré ball is one model of hyperbolic space (constant negative curvature). Other models: + +- **Hyperboloid**: Upper sheet of hyperboloid +- **Klein**: Projective model +- **Upper half-space**: Complex plane with Im(z) > 0 + +All are isometric (same geometry, different coordinates). + +### Geodesics + +Shortest paths in Poincaré ball are: + +- **Through origin**: Straight lines +- **Not through origin**: Circular arcs perpendicular to boundary + +### Curvature + +Poincaré ball has constant negative curvature κ = -1. This creates exponential growth in volume: + +``` +Vol(ball of radius r) ∝ exp(r) +``` + +This is why it's perfect for trees! + +## Further Reading + +### Papers + +- [Poincaré Embeddings (Nickel & Kiela, 2017)](https://arxiv.org/abs/1705.08039) +- [Hyperbolic Neural Networks (Ganea et al., 2018)](https://arxiv.org/abs/1805.09112) +- [Learning Continuous Hierarchies (Sala et al., 2018)](https://arxiv.org/abs/1806.03417) + +### Books + +- Anderson, J. W. (2005). Hyperbolic Geometry (Springer Undergraduate Mathematics) +- Ratcliffe, J. G. (2006). Foundations of Hyperbolic Manifolds + +### Implementations + +- [geoopt](https://github.com/geoopt/geoopt) - Riemannian optimization in PyTorch +- [hyperlib](https://github.com/lateral/hyperlib) - Hyperbolic geometry utilities diff --git a/docs/user-guide.md b/docs/user-guide.md new file mode 100644 index 0000000..4f2ce11 --- /dev/null +++ b/docs/user-guide.md @@ -0,0 +1,401 @@ +# User Guide + +Comprehensive guide to using taxembed for hierarchical taxonomy embeddings. + +## Installation + +### Using uv (recommended) + +```bash +git clone https://github.com/yourusername/taxembed.git +cd taxembed +uv sync +``` + +### Using pip + +```bash +pip install -e . +``` + +## Quick Start + +### 1. Download NCBI Taxonomy + +```bash +taxembed-download +``` + +This downloads and extracts: + +- `data/nodes.dmp` - Taxonomy structure +- `data/names.dmp` - Organism names +- `data/taxonomy_edges.edgelist` - Parent-child relationships + +### 2. Prepare Training Data + +```bash +taxembed-prepare +``` + +This builds the transitive closure (all ancestor-descendant pairs): + +- Input: `data/taxonomy_edges_small.edgelist` +- Output: `data/taxonomy_edges_small_transitive.pkl` (975K training pairs) + +### 3. Train Model + +```bash +taxembed-train +``` + +Or with custom parameters: + +```bash +taxembed-train --epochs 100 --dim 10 --lambda-reg 0.1 +``` + +Training options: + +- `--epochs`: Number of training epochs (default: 100) +- `--dim`: Embedding dimensionality (default: 10) +- `--batch-size`: Batch size (default: 64) +- `--lr`: Learning rate (default: 0.005) +- `--margin`: Ranking loss margin (default: 0.2) +- `--lambda-reg`: Regularization strength (default: 0.1) +- `--early-stopping`: Patience for early stopping (default: 5) + +### 4. Visualize Results + +```bash +taxembed-visualize taxonomy_model_small_best.pth +``` + +## CLI Commands + +### Unified `taxembed` Command + +The main entry point supports multiple subcommands: + +```bash +# Train any clade by name or TaxID +taxembed train Cnidaria -as cnidaria --epochs 100 + +# Visualize trained model +taxembed visualize cnidaria --children 1 + +# Analyze hierarchy quality +taxembed analyze cnidaria_best.pth +``` + +### Legacy Commands + +Individual commands are still available: + +- `taxembed-download` - Download NCBI taxonomy +- `taxembed-prepare` - Build transitive closure +- `taxembed-train` - Train embeddings +- `taxembed-visualize` - Create UMAP visualizations +- `taxembed-check` - Run sanity checks + +## Working with Custom Clades + +### Build Custom Dataset + +```python +from taxembed.builders import build_clade_dataset + +result = build_clade_dataset( + root_taxid=33208, # Metazoa (animals) + dataset_name="animals", + output_dir="data/taxopy/animals" +) + +print(f"Created dataset with {result.node_count:,} nodes") +``` + +### Train on Custom Data + +```bash +taxembed train 33208 -as animals --epochs 100 +``` + +Or using Python: + +```python +from taxembed.cli.train import main +import sys + +sys.argv = [ + 'train', + '--data', 'data/taxopy/animals/taxonomy_edges_animals_transitive.pkl', + '--mapping', 'data/taxopy/animals/taxonomy_edges_animals.mapping.tsv', + '--checkpoint', 'animals_model.pth', + '--epochs', '100' +] + +main() +``` + +## Understanding the Model + +### Poincaré Ball Model + +Taxembed uses hyperbolic geometry (Poincaré ball model) to represent hierarchies: + +- **Center**: Root of taxonomy (Cellular organisms) +- **Boundary**: Leaf nodes (species/strains) +- **Distance from center**: Depth in hierarchy +- **Angular distance**: Similarity within level + +### Training Features + +1. **Transitive Closure**: Trains on ALL ancestor-descendant pairs, not just parent-child +2. **Depth-Aware Initialization**: Deeper nodes start closer to boundary +3. **Radial Regularization**: Encourages ||embedding|| ≈ f(depth) +4. **Hard Negative Sampling**: Samples cousins at same depth +5. **Depth Weighting**: Deeper pairs weighted more heavily + +### Ball Constraint Enforcement + +Three-layer strategy ensures 100% valid embeddings: + +1. **Gradient clipping**: Prevents large updates +2. **Selective projection**: Projects updated embeddings after each batch +3. **Full projection**: Projects all embeddings at epoch end + +## Advanced Usage + +### Custom Training Loop + +```python +import torch +from taxembed.models import HierarchicalPoincareEmbedding +from taxembed.training import HierarchicalDataLoader, train_model + +# Load data +with open('data/training_data.pkl', 'rb') as f: + training_data = pickle.load(f) + +# Create model +model = HierarchicalPoincareEmbedding( + n_nodes=10000, + dim=10, + max_depth=38, + init_depth_data=idx_to_depth +) + +# Create data loader +dataloader = HierarchicalDataLoader( + training_data=training_data, + n_nodes=10000, + batch_size=64, + n_negatives=50 +) + +# Train +optimizer = torch.optim.Adam(model.parameters(), lr=0.005) +train_model( + model, dataloader, optimizer, + n_epochs=100, + idx_to_depth=idx_to_depth, + max_depth=38, + device=torch.device('cpu'), + checkpoint_base='my_model.pth' +) +``` + +### Loading Trained Embeddings + +```python +import torch +import pandas as pd + +# Load checkpoint +ckpt = torch.load('taxonomy_model_small_best.pth') +embeddings = ckpt['embeddings'] # Shape: (n_nodes, dim) + +# Load TaxID mapping +mapping = pd.read_csv('data/taxonomy_edges_small.mapping.tsv', + sep='\t', header=None, names=['idx', 'taxid']) + +# Get embedding for specific TaxID +taxid = 9606 # Homo sapiens +idx = mapping[mapping['taxid'] == str(taxid)]['idx'].iloc[0] +human_embedding = embeddings[idx] +``` + +### Nearest Neighbors + +```python +import torch + +def find_nearest_neighbors(query_idx, embeddings, model, k=10): + """Find k nearest neighbors in hyperbolic space.""" + query_emb = embeddings[query_idx].unsqueeze(0) + all_embs = embeddings + + # Compute Poincaré distances + distances = model.poincare_distance( + query_emb.expand(len(embeddings), -1), + all_embs + ) + + # Get top k + _, indices = torch.topk(distances, k, largest=False) + return indices, distances[indices] +``` + +## Troubleshooting + +### Out of Memory + +Reduce batch size: + +```bash +taxembed-train --batch-size 32 +``` + +### Training Too Slow + +Increase batch size and reduce negatives: + +```bash +taxembed-train --batch-size 128 --n-negatives 25 +``` + +### Poor Hierarchy Quality + +- Train longer (100+ epochs) +- Increase regularization: `--lambda-reg 0.2` +- Try larger embedding dimension: `--dim 20` + +### Embeddings Outside Ball + +This shouldn't happen with current implementation. If it does: + +- Check for NaN/Inf in data +- Reduce learning rate: `--lr 0.001` +- Increase regularization + +## Best Practices + +1. **Start with small dataset**: Test on `taxonomy_edges_small` first +2. **Monitor metrics**: Watch for decreasing loss and stable norms +3. **Use early stopping**: Prevents overfitting (default: 5 epochs patience) +4. **Save checkpoints**: Models save best checkpoint automatically +5. **Validate results**: Use `taxembed-check` to verify data quality + +## Performance + +### Small Dataset (111K nodes) + +- Training time: ~3 min/epoch on M3 Mac CPU +- Recommended epochs: 50-100 +- Expected loss: 0.47 after 28 epochs + +### Full Dataset (2.7M nodes) + +- Training time: ~60 min/epoch on M3 Mac CPU +- Recommended epochs: 20-50 +- Memory: ~8GB RAM + +## Development + +### Code Quality Tools + +The project uses modern Python tooling to maintain high code quality: + +#### Linting with Ruff + +Ruff provides fast linting and formatting: + +```bash +# Check for linting issues +uv run ruff check . + +# Auto-fix safe issues +uv run ruff check --fix . + +# Auto-fix including unsafe fixes (e.g., remove unused imports) +uv run ruff check --fix --unsafe-fixes . + +# Show detailed error messages +uv run ruff check --output-format=full . +``` + +#### Formatting with Ruff + +Ensure consistent code style: + +```bash +# Format all Python files +uv run ruff format . + +# Check formatting without applying changes +uv run ruff format --check . +``` + +#### Type Checking with MyPy + +Static type analysis catches bugs early: + +```bash +# Check all source code +uv run mypy src/taxembed + +# Check specific module +uv run mypy src/taxembed/models/ + +# Show error codes for better understanding +uv run mypy --show-error-codes src/taxembed +``` + +The project uses **gradual typing**, meaning type hints are being added incrementally. Most modules are currently exempt from strict type checking (see `pyproject.toml`). + +#### Testing with Pytest + +Run the test suite: + +```bash +# Run all tests +uv run pytest + +# Run with coverage report +uv run pytest --cov=src/taxembed --cov-report=term-missing + +# Run specific test file +uv run pytest tests/test_models.py + +# Run with verbose output +uv run pytest -v +``` + +#### Complete Quality Check + +Run all checks at once before committing: + +```bash +uv run ruff check . && \ +uv run ruff format --check . && \ +uv run mypy src/taxembed && \ +uv run pytest +``` + +### Configuration + +All tools are configured in `pyproject.toml`: + +- **Ruff**: Line length 100, Python 3.11+, comprehensive rule set +- **MyPy**: Strict typing with gradual adoption +- **Pytest**: Coverage reporting enabled + +See [CONTRIBUTING.md](../CONTRIBUTING.md) for detailed development guidelines. + +--- + +## References + +- [Poincaré Embeddings Paper](https://arxiv.org/abs/1705.08039) +- [NCBI Taxonomy](https://www.ncbi.nlm.nih.gov/taxonomy) +- [Examples](../examples/) diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000..e353d35 --- /dev/null +++ b/examples/README.md @@ -0,0 +1,56 @@ +# Examples + +This directory contains example scripts demonstrating how to use the taxembed package. + +## Available Examples + +### basic_training.py + +Simple example showing how to train a Poincaré embedding model on taxonomy data. + +```bash +uv run python examples/basic_training.py +``` + +### custom_dataset.py + +Example of using taxembed with custom taxonomy data. + +```bash +uv run python examples/custom_dataset.py --root-taxid 33208 --name animals +``` + +### visualize_groups.py + +Demonstrates how to create visualizations highlighting specific taxonomic groups. + +```bash +uv run python examples/visualize_groups.py --checkpoint model.pth +``` + +### nn_demo.py + +Interactive demo for finding nearest neighbors in the embedding space. + +```bash +uv run python examples/nn_demo.py +``` + +## Running Examples + +All examples can be run using `uv`: + +```bash +# Install dependencies first +uv sync + +# Run any example +uv run python examples/.py +``` + +## Learning Path + +1. Start with `basic_training.py` to understand the training workflow +2. Try `custom_dataset.py` to work with specific taxonomic clades +3. Use `visualize_groups.py` to analyze your trained models +4. Explore `nn_demo.py` for interactive analysis diff --git a/examples/basic_training.py b/examples/basic_training.py new file mode 100644 index 0000000..2a418c2 --- /dev/null +++ b/examples/basic_training.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Basic example of training Poincaré embeddings on taxonomy data.""" + +import pickle +from pathlib import Path + +import pandas as pd +import torch +import torch.optim as optim + +from taxembed.models import HierarchicalPoincareEmbedding +from taxembed.training import HierarchicalDataLoader, train_model + + +def main(): + """Train a simple model on the small taxonomy dataset.""" + + # Paths + data_dir = Path("data") + training_data_file = data_dir / "taxonomy_edges_small_transitive.pkl" + mapping_file = data_dir / "taxonomy_edges_small.mapping.tsv" + + # Check if data exists + if not training_data_file.exists(): + print(f"❌ Training data not found: {training_data_file}") + print(" Run: taxembed-download && taxembed-prepare") + return + + print("=" * 60) + print("BASIC TRAINING EXAMPLE") + print("=" * 60) + + # Load training data + print("\n1. Loading training data...") + with open(training_data_file, "rb") as f: + training_data = pickle.load(f) + print(f" ✓ Loaded {len(training_data):,} training pairs") + + # Load mapping + print("\n2. Loading node mapping...") + mapping_df = pd.read_csv(mapping_file, sep="\t", header=None, names=["idx", "taxid"]) + n_nodes = len(mapping_df) + print(f" ✓ {n_nodes:,} unique nodes") + + # Build depth map + print("\n3. Building depth information...") + idx_to_depth = {} + max_depth = 0 + for item in training_data: + idx_to_depth[item["ancestor_idx"]] = item["ancestor_depth"] + idx_to_depth[item["descendant_idx"]] = item["descendant_depth"] + max_depth = max(max_depth, item["ancestor_depth"], item["descendant_depth"]) + print(f" ✓ Depth range: [0, {max_depth}]") + + # Create model + print("\n4. Creating Poincaré embedding model...") + model = HierarchicalPoincareEmbedding( + n_nodes=n_nodes, + dim=10, # 10-dimensional embeddings + max_depth=max_depth, + init_depth_data=idx_to_depth, + ) + print(f" ✓ Model initialized with {n_nodes:,} nodes in 10D") + + # Create data loader + print("\n5. Creating data loader...") + dataloader = HierarchicalDataLoader( + training_data=training_data, + n_nodes=n_nodes, + batch_size=64, + n_negatives=50, + depth_stratify=True, + ) + + # Optimizer + optimizer = optim.Adam(model.parameters(), lr=0.005) + + # Train for a few epochs (just for demo) + print("\n6. Training model...") + print(" (Training for 5 epochs as a demo - increase for better results)\n") + + device = torch.device("cpu") # Use CPU for demo + + train_model( + model=model, + dataloader=dataloader, + optimizer=optimizer, + n_epochs=5, # Short for demo + idx_to_depth=idx_to_depth, + max_depth=max_depth, + device=device, + margin=0.2, + lambda_reg=0.1, + early_stopping_patience=0, # Disabled for demo + checkpoint_base="examples/demo_model.pth", + ) + + print("\n" + "=" * 60) + print("TRAINING COMPLETE!") + print("=" * 60) + print("\nNext steps:") + print(" - Visualize: taxembed-visualize examples/demo_model_best.pth") + print(" - Train longer: Increase n_epochs for better quality") + print(" - Analyze: python examples/nn_demo.py") + + +if __name__ == "__main__": + main() diff --git a/examples/custom_dataset.py b/examples/custom_dataset.py new file mode 100644 index 0000000..5af0f65 --- /dev/null +++ b/examples/custom_dataset.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +"""Example of building and training on a custom taxonomic clade.""" + +import argparse +from pathlib import Path + +from taxembed.builders import build_clade_dataset + + +def main(): + """Build a custom dataset and show how to use it.""" + parser = argparse.ArgumentParser(description="Build custom taxonomy dataset") + parser.add_argument( + "--root-taxid", + type=int, + required=True, + help="Root TaxID for the clade (e.g., 33208 for Metazoa)", + ) + parser.add_argument("--name", required=True, help="Dataset name (e.g., animals, plants)") + parser.add_argument( + "--max-depth", + type=int, + default=None, + help="Maximum depth to include (optional)", + ) + + args = parser.parse_args() + + print("=" * 60) + print("CUSTOM DATASET BUILDER") + print("=" * 60) + + # Build the dataset + print(f"\nBuilding dataset for TaxID {args.root_taxid} ({args.name})...") + + data_dir = Path("data") + output_dir = data_dir / "taxopy" / args.name + + result = build_clade_dataset( + root_taxid=args.root_taxid, + dataset_name=args.name, + output_dir=output_dir, + taxdump_dir=data_dir, + max_depth=args.max_depth, + ) + + print("\n" + "=" * 60) + print("DATASET READY!") + print("=" * 60) + print("\nStatistics:") + print(f" Root TaxID: {result.root_taxid}") + print(f" Nodes: {result.node_count:,}") + print(f" Training pairs: {result.pairs_count:,}") + print(f" Max depth: {result.max_depth}") + print("\nFiles created:") + for name, path in result.files.items(): + print(f" {name}: {path}") + + print("\nNext steps:") + print(f" Train: taxembed train {args.root_taxid} -as {args.name} --epochs 100") + print(" Or: python -c 'from taxembed.cli.train import main; main()' \\") + print(f" --data {result.files['transitive_pickle']} \\") + print(f" --mapping {result.files['mapping']} \\") + print(f" --checkpoint {args.name}_model.pth") + + +if __name__ == "__main__": + main() diff --git a/nn_demo.py b/examples/nn_demo.py similarity index 79% rename from nn_demo.py rename to examples/nn_demo.py index 3be4329..5b5c19c 100644 --- a/nn_demo.py +++ b/examples/nn_demo.py @@ -1,4 +1,8 @@ -import sys, torch, numpy as np, pandas as pd +import sys + +import numpy as np +import pandas as pd +import torch ckpt_path, map_path, query = sys.argv[1], sys.argv[2], sys.argv[3] ckpt = torch.load(ckpt_path, map_location="cpu") @@ -15,8 +19,8 @@ # Load mapping m = pd.read_csv(map_path, sep="\t") -tax2idx = dict(zip(m["taxid"].astype(str), m["idx"])) -idx2tax = dict(zip(m["idx"], m["taxid"].astype(str))) +tax2idx = dict(zip(m["taxid"].astype(str), m["idx"], strict=False)) +idx2tax = dict(zip(m["idx"], m["taxid"].astype(str), strict=False)) # Query i = tax2idx[str(query)] diff --git a/examples/visualize_groups.py b/examples/visualize_groups.py new file mode 100644 index 0000000..ff97aa5 --- /dev/null +++ b/examples/visualize_groups.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 +"""Example of visualizing embeddings with custom group highlighting.""" + +import argparse + +from taxembed.visualization import ( + load_embeddings, + load_mapping, +) + + +def main(): + """Create custom visualizations of trained embeddings.""" + parser = argparse.ArgumentParser(description="Visualize embeddings") + parser.add_argument("checkpoint", help="Path to checkpoint file") + parser.add_argument( + "--mapping", + default="data/taxonomy_edges_small.mapping.tsv", + help="Path to mapping file", + ) + parser.add_argument("--output", default="embedding_viz.png", help="Output image path") + parser.add_argument("--sample", type=int, default=10000, help="Number of points to visualize") + + args = parser.parse_args() + + print("=" * 60) + print("EMBEDDING VISUALIZATION") + print("=" * 60) + + # Load embeddings + print(f"\nLoading embeddings from {args.checkpoint}...") + embeddings = load_embeddings(args.checkpoint) + print(f" ✓ Loaded {embeddings.shape[0]:,} embeddings of dimension {embeddings.shape[1]}") + + # Load mapping + print(f"\nLoading mapping from {args.mapping}...") + tax2idx, idx2tax = load_mapping(args.mapping) + if tax2idx: + print(f" ✓ Loaded {len(tax2idx):,} TaxID mappings") + + # Create visualization + print(f"\nCreating UMAP visualization (sampling {args.sample:,} points)...") + print(" This may take a few minutes...") + + # Note: The full create_umap_visualization function needs to be properly + # extracted from visualize_multi_groups.py. For now, this is a placeholder. + print("\n ⚠️ Full visualization function needs to be refactored.") + print(f" For now, use: taxembed-visualize {args.checkpoint}") + + print("\n" + "=" * 60) + print("TIP: Use the taxembed CLI for full visualization features:") + print(f" taxembed visualize --sample {args.sample}") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 204a700..b6b81f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,15 +1,43 @@ +# ============================================================================= +# Build System & Configuration +# ============================================================================= [build-system] requires = ["hatchling"] build-backend = "hatchling.build" +[tool.hatch.build.targets.wheel] +packages = ["src/taxembed"] + +# ============================================================================= +# Project Metadata +# ============================================================================= [project] name = "taxembed" -version = "0.2.0" -description = "Hierarchical Poincaré embeddings for NCBI biological taxonomy" +version = "1.0.0" +description = "Hierarchical Poincaré embeddings for biological taxonomy" readme = "README.md" license = { text = "MIT" } authors = [ - { name = "Joha Coludar", email = "jcoludar@gmail.com" } + { name = "Thomas Senoner", email = "info@example.com" } +] +keywords = [ + "taxonomy", + "embeddings", + "hyperbolic", + "poincare", + "hierarchical", + "ncbi", + "biology", +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering :: Bio-Informatics", + "Topic :: Scientific/Engineering :: Artificial Intelligence", ] requires-python = ">=3.11" dependencies = [ @@ -23,14 +51,6 @@ dependencies = [ "taxopy>=0.14.0", ] -[project.scripts] -taxembed = "taxembed.cli.main:main" -taxembed-download = "taxembed.cli.download:main" -taxembed-prepare = "taxembed.cli.prepare:main" -taxembed-train = "taxembed.cli.train:main" -taxembed-visualize = "taxembed.cli.visualize:main" -taxembed-check = "taxembed.cli.check:main" - [project.optional-dependencies] dev = [ "ruff>=0.6.0", @@ -39,22 +59,64 @@ dev = [ "mypy>=1.0.0", ] +[project.scripts] +taxembed = "taxembed.cli.main:main" +taxembed-download = "taxembed.cli.download:main" +taxembed-prepare = "taxembed.cli.prepare:main" +taxembed-train = "taxembed.cli.train:main" +taxembed-visualize = "taxembed.cli.visualize:main" +taxembed-check = "taxembed.cli.check:main" +taxembed-analyze = "taxembed.cli.analyze:main" + [project.urls] Repository = "https://github.com/jcoludar/taxembed" Documentation = "https://github.com/jcoludar/taxembed/blob/main/README.md" Issues = "https://github.com/jcoludar/taxembed/issues" -[dependency-groups] -dev = [ - "ruff>=0.6.0", - "pytest>=8.0.0", - "pytest-cov>=4.1.0", - "mypy>=1.0.0", +# ============================================================================= +# Development Tools Configuration +# ============================================================================= + +# ----------------------------------------------------------------------------- +# MyPy: Static Type Checking +# ----------------------------------------------------------------------------- +[tool.mypy] +python_version = "3.11" +warn_return_any = true +warn_unused_configs = true +disallow_untyped_defs = true +strict_optional = true +check_untyped_defs = true +ignore_missing_imports = true # Allow untyped third-party libraries + +# Gradual typing: temporarily allow untyped defs in these modules +# Remove modules from this list as you add complete type annotations +[[tool.mypy.overrides]] +module = [ + "taxembed.analysis.*", + "taxembed.builders.*", + "taxembed.cli.*", + "taxembed.data.*", + "taxembed.models.*", + "taxembed.training.*", + "taxembed.validation.*", + "taxembed.visualization.*", ] +disallow_untyped_defs = false +check_untyped_defs = false +warn_return_any = false -[tool.hatch.build.targets.wheel] -packages = ["src/taxembed"] +# ----------------------------------------------------------------------------- +# Pytest: Testing Framework +# ----------------------------------------------------------------------------- +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +addopts = "--cov=src/taxembed --cov-report=term-missing" +# ----------------------------------------------------------------------------- +# Ruff: Linting and Formatting +# ----------------------------------------------------------------------------- [tool.ruff] line-length = 100 target-version = "py311" @@ -68,7 +130,8 @@ exclude = [ "__pycache__", ".pytest_cache", ".ruff_cache", - "hype", # Original Facebook code + "*.egg-info", + "_vendor", # Facebook's original code (backup) ] [tool.ruff.lint] @@ -83,24 +146,19 @@ select = [ ] ignore = [ "E501", # line too long (handled by formatter) - "W503", # line break before binary operator + "E203", # whitespace before ':' "B008", # function calls in argument defaults + "C901", # function is too complex (acceptable for now) ] [tool.ruff.lint.isort] known-first-party = ["taxembed"] +section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["F401", "F403"] +"tests/*" = ["F841"] [tool.ruff.format] quote-style = "double" indent-style = "space" - -[tool.pytest.ini_options] -testpaths = ["tests"] -python_files = ["test_*.py"] -addopts = "--cov=src/taxembed --cov-report=term-missing" - -[tool.mypy] -python_version = "3.11" -warn_return_any = true -warn_unused_configs = true -disallow_untyped_defs = false # Gradual typing diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 0a4304f..0000000 --- a/requirements.txt +++ /dev/null @@ -1,10 +0,0 @@ -iopath -nltk -scikit-learn -pandas -h5py -cython -tqdm -numpy>=1.21.0,<2.0 -torch>=2.0.0 -taxopy>=0.14.0 diff --git a/ruff.toml b/ruff.toml deleted file mode 100644 index 9e7d8b3..0000000 --- a/ruff.toml +++ /dev/null @@ -1,42 +0,0 @@ -# Ruff configuration for taxembed project -# See https://docs.astral.sh/ruff/configuration/ - -line-length = 100 -target-version = "py38" - -exclude = [ - ".git", - ".venv", - "venv", - "venv311", - "build", - "dist", - "__pycache__", - ".pytest_cache", - ".ruff_cache", - "*.egg-info", -] - -[lint] -select = [ - "E", # pycodestyle errors - "W", # pycodestyle warnings - "F", # pyflakes - "I", # isort - "C", # flake8-comprehensions - "B", # flake8-bugbear - "UP", # pyupgrade -] -ignore = [ - "E501", # line too long (handled by formatter) - "W503", # line break before binary operator - "E203", # whitespace before ':' -] - -[lint.isort] -known-first-party = ["taxembed"] -section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"] - -[lint.per-file-ignores] -"__init__.py" = ["F401", "F403"] -"tests/*" = ["F841"] diff --git a/sanity_check.py b/sanity_check.py deleted file mode 100644 index a69de93..0000000 --- a/sanity_check.py +++ /dev/null @@ -1,412 +0,0 @@ -#!/usr/bin/env python3 -""" -Comprehensive sanity check for hierarchical training pipeline. -Verifies data integrity, model logic, and training setup. -""" - -import torch -import numpy as np -import pickle -import pandas as pd -from collections import defaultdict - -print("="*80) -print("COMPREHENSIVE SANITY CHECK") -print("="*80) -print() - -# ============================================================================ -# 1. MAPPING FILE INTEGRITY -# ============================================================================ -print("1. MAPPING FILE INTEGRITY") -print("-" * 40) - -mapping_file = "data/taxonomy_edges_small.mapping.tsv" -df = pd.read_csv(mapping_file, sep="\t") - -print(f" Columns: {list(df.columns)}") -print(f" Shape: {df.shape}") -print(f" First 3 rows:") -print(df.head(3)) - -# Check for duplicates -dup_taxids = df[df.duplicated(subset=['taxid'], keep=False)] -dup_indices = df[df.duplicated(subset=['idx'], keep=False)] - -if len(dup_taxids) > 0: - print(f" ❌ ERROR: {len(dup_taxids)} duplicate TaxIDs found!") -else: - print(f" ✅ No duplicate TaxIDs") - -if len(dup_indices) > 0: - print(f" ❌ ERROR: {len(dup_indices)} duplicate indices found!") -else: - print(f" ✅ No duplicate indices") - -# Check index continuity -indices = sorted(df['idx'].values) -expected_indices = list(range(len(df))) -if indices != expected_indices: - print(f" ❌ ERROR: Index discontinuity!") - print(f" Expected: 0-{len(df)-1}") - print(f" Got: {min(indices)}-{max(indices)}") -else: - print(f" ✅ Indices are continuous: 0-{max(indices)}") - -print() - -# ============================================================================ -# 2. TRANSITIVE CLOSURE DATA -# ============================================================================ -print("2. TRANSITIVE CLOSURE DATA") -print("-" * 40) - -with open("data/taxonomy_edges_small_transitive.pkl", "rb") as f: - training_data = pickle.load(f) - -print(f" Total pairs: {len(training_data):,}") - -# Check indices are in valid range -max_mapping_idx = df['idx'].max() -all_indices = [] -for item in training_data: - all_indices.append(item['ancestor_idx']) - all_indices.append(item['descendant_idx']) - -max_idx = max(all_indices) -min_idx = min(all_indices) - -print(f" Index range in data: {min_idx} - {max_idx}") -print(f" Mapping index range: 0 - {max_mapping_idx}") - -if max_idx > max_mapping_idx: - print(f" ❌ ERROR: Data has indices ({max_idx}) > mapping max ({max_mapping_idx})!") -else: - print(f" ✅ All indices within mapping range") - -# Check for self-loops -self_loops = sum(1 for item in training_data if item['ancestor_idx'] == item['descendant_idx']) -if self_loops > 0: - print(f" ❌ WARNING: {self_loops} self-loops found!") -else: - print(f" ✅ No self-loops") - -# Check depth consistency -invalid_depths = [] -for i, item in enumerate(training_data[:1000]): # Sample check - expected_diff = item['descendant_depth'] - item['ancestor_depth'] - if item['depth_diff'] != expected_diff: - invalid_depths.append(i) - -if invalid_depths: - print(f" ❌ ERROR: {len(invalid_depths)} pairs with invalid depth_diff in sample!") -else: - print(f" ✅ Depth differences are consistent") - -# Check depth values are non-negative -negative_depths = sum(1 for item in training_data if item['depth_diff'] <= 0) -if negative_depths > 0: - print(f" ❌ ERROR: {negative_depths} pairs with non-positive depth_diff!") -else: - print(f" ✅ All depth differences are positive") - -print() - -# ============================================================================ -# 3. PROJECTION LOGIC TEST -# ============================================================================ -print("3. PROJECTION LOGIC TEST") -print("-" * 40) - -def test_projection(n=100, dim=10): - """Test that projection correctly constrains embeddings to unit ball.""" - - # Create random embeddings (some outside ball) - embeddings = torch.randn(n, dim) * 2.0 # Scale up to force some outside - - # Count how many are outside ball before projection - norms_before = embeddings.norm(dim=1) - outside_before = (norms_before >= 1.0).sum().item() - - # Project - eps = 1e-5 - norms = embeddings.norm(dim=1, keepdim=True) - scale = torch.clamp(norms, max=1 - eps) / (norms + eps) - embeddings_projected = embeddings * scale - - # Check after projection - norms_after = embeddings_projected.norm(dim=1) - outside_after = (norms_after >= 1.0).sum().item() - max_norm = norms_after.max().item() - - print(f" Before projection: {outside_before}/{n} embeddings outside ball") - print(f" After projection: {outside_after}/{n} embeddings outside ball") - print(f" Max norm after projection: {max_norm:.6f}") - print(f" Target max norm: {1 - eps:.6f}") - - if outside_after > 0: - print(f" ❌ ERROR: Projection failed! {outside_after} embeddings still outside!") - return False - elif max_norm > 1.0: - print(f" ❌ ERROR: Max norm {max_norm} > 1.0!") - return False - else: - print(f" ✅ Projection working correctly") - return True - -test_projection() -print() - -# ============================================================================ -# 4. HYPERBOLIC DISTANCE TEST -# ============================================================================ -print("4. HYPERBOLIC DISTANCE TEST") -print("-" * 40) - -def test_poincare_distance(): - """Test Poincaré distance computation.""" - - eps = 1e-5 - - # Test case 1: Distance to self should be 0 - u = torch.tensor([[0.1, 0.2, 0.0]], dtype=torch.float32) - v = u.clone() - - u_norm_sq = (u ** 2).sum(dim=-1) - v_norm_sq = (v ** 2).sum(dim=-1) - u_norm_sq = torch.clamp(u_norm_sq, 0, 1 - eps) - v_norm_sq = torch.clamp(v_norm_sq, 0, 1 - eps) - diff_norm_sq = ((u - v) ** 2).sum(dim=-1) - numerator = 2 * diff_norm_sq - denominator = (1 - u_norm_sq) * (1 - v_norm_sq) - dist = torch.acosh(1 + numerator / (denominator + eps) + eps) - - print(f" Test 1 - Distance to self:") - print(f" Distance: {dist.item():.6f}") - if dist.item() < 0.01: - print(f" ✅ Correct (≈0)") - else: - print(f" ❌ ERROR: Should be ≈0") - - # Test case 2: Distance increases with separation - u = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32) - v1 = torch.tensor([[0.1, 0.0, 0.0]], dtype=torch.float32) - v2 = torch.tensor([[0.5, 0.0, 0.0]], dtype=torch.float32) - - def compute_dist(a, b): - a_norm_sq = (a ** 2).sum(dim=-1) - b_norm_sq = (b ** 2).sum(dim=-1) - a_norm_sq = torch.clamp(a_norm_sq, 0, 1 - eps) - b_norm_sq = torch.clamp(b_norm_sq, 0, 1 - eps) - diff_norm_sq = ((a - b) ** 2).sum(dim=-1) - numerator = 2 * diff_norm_sq - denominator = (1 - a_norm_sq) * (1 - b_norm_sq) - return torch.acosh(1 + numerator / (denominator + eps) + eps) - - dist1 = compute_dist(u, v1) - dist2 = compute_dist(u, v2) - - print(f" Test 2 - Distance monotonicity:") - print(f" d(origin, 0.1) = {dist1.item():.6f}") - print(f" d(origin, 0.5) = {dist2.item():.6f}") - if dist2 > dist1: - print(f" ✅ Correct (larger separation = larger distance)") - else: - print(f" ❌ ERROR: Distance should increase with separation") - -test_poincare_distance() -print() - -# ============================================================================ -# 5. DEPTH-AWARE INITIALIZATION TEST -# ============================================================================ -print("5. DEPTH-AWARE INITIALIZATION TEST") -print("-" * 40) - -def test_depth_initialization(): - """Test that depth-aware initialization works correctly.""" - - max_depth = 38 - - # Test different depths - test_depths = [0, 10, 20, 30, 38] - - print(f" Expected radius by depth:") - for depth in test_depths: - target_radius = 0.1 + (depth / max_depth) * 0.85 - print(f" Depth {depth:2d}: r = {target_radius:.4f}") - - print() - print(f" Properties:") - r_min = 0.1 + (0 / max_depth) * 0.85 - r_max = 0.1 + (max_depth / max_depth) * 0.85 - print(f" Root (depth 0): r ≈ {r_min:.4f}") - print(f" Leaves (depth {max_depth}): r ≈ {r_max:.4f}") - print(f" All radii < 1.0: {r_max < 1.0}") - - if r_max < 1.0: - print(f" ✅ All initialized embeddings will be inside ball") - else: - print(f" ❌ ERROR: Max radius {r_max} >= 1.0!") - - return r_max - -r_max_global = test_depth_initialization() -print() - -# ============================================================================ -# 6. SIBLING MAP LOGIC TEST -# ============================================================================ -print("6. SIBLING MAP LOGIC TEST") -print("-" * 40) - -# Build sibling map -sibling_map = defaultdict(list) -depth_buckets = defaultdict(list) - -for item in training_data: - desc_idx = item['descendant_idx'] - desc_depth = item['descendant_depth'] - depth_buckets[desc_depth].append(desc_idx) - -for depth, nodes in depth_buckets.items(): - for node in nodes: - # Siblings are other nodes at same depth (excluding self) - siblings = [n for n in nodes if n != node] - sibling_map[node] = siblings - -# Check sibling map -total_nodes_with_siblings = len(sibling_map) -nodes_with_no_siblings = sum(1 for siblings in sibling_map.values() if len(siblings) == 0) -avg_siblings = np.mean([len(s) for s in sibling_map.values()]) - -print(f" Nodes with sibling info: {total_nodes_with_siblings:,}") -print(f" Nodes with no siblings: {nodes_with_no_siblings:,}") -print(f" Average siblings per node: {avg_siblings:.1f}") - -if nodes_with_no_siblings > total_nodes_with_siblings * 0.5: - print(f" ⚠️ WARNING: Many nodes have no siblings (hard negatives will fallback to random)") -else: - print(f" ✅ Most nodes have siblings for hard negative sampling") - -# Sample check: verify siblings are actually at same depth -sample_node = list(sibling_map.keys())[0] -sample_siblings = sibling_map[sample_node] -node_depth = None -for item in training_data: - if item['descendant_idx'] == sample_node: - node_depth = item['descendant_depth'] - break - -if node_depth is not None and len(sample_siblings) > 0: - sibling_depths = set() - for item in training_data: - if item['descendant_idx'] in sample_siblings[:10]: # Check first 10 - sibling_depths.add(item['descendant_depth']) - - if len(sibling_depths) == 1 and node_depth in sibling_depths: - print(f" ✅ Siblings are at same depth (verified sample)") - else: - print(f" ❌ ERROR: Siblings have different depths!") - -print() - -# ============================================================================ -# 7. REGULARIZER TARGET CHECK -# ============================================================================ -print("7. REGULARIZER TARGET CHECK") -print("-" * 40) - -# Build idx_to_depth from training data -idx_to_depth = {} -for item in training_data: - idx_to_depth[item['descendant_idx']] = item['descendant_depth'] - if item['ancestor_idx'] not in idx_to_depth: - idx_to_depth[item['ancestor_idx']] = item['ancestor_depth'] - -max_depth = max(idx_to_depth.values()) -n_nodes = max(max(item['ancestor_idx'], item['descendant_idx']) for item in training_data) + 1 - -print(f" Nodes in training: {n_nodes:,}") -print(f" Nodes with depth info: {len(idx_to_depth):,}") -print(f" Max depth: {max_depth}") - -# Check that all regularized nodes have valid depth -nodes_without_depth = n_nodes - len(idx_to_depth) -if nodes_without_depth > 0: - print(f" ⚠️ {nodes_without_depth:,} nodes have no depth info (won't be regularized)") -else: - print(f" ✅ All nodes have depth info") - -# Check regularizer targets are valid -invalid_targets = 0 -for idx, depth in list(idx_to_depth.items())[:1000]: # Sample check - target_radius = 0.1 + (depth / max_depth) * 0.85 - if target_radius >= 1.0: - invalid_targets += 1 - -if invalid_targets > 0: - print(f" ❌ ERROR: {invalid_targets} regularizer targets >= 1.0!") -else: - print(f" ✅ All regularizer targets < 1.0") - -print() - -# ============================================================================ -# 8. BATCH SIZE VS DATASET SIZE -# ============================================================================ -print("8. TRAINING CONFIGURATION") -print("-" * 40) - -batch_size = 64 -n_training_pairs = len(training_data) -n_batches_per_epoch = (n_training_pairs + batch_size - 1) // batch_size - -print(f" Training pairs: {n_training_pairs:,}") -print(f" Batch size: {batch_size}") -print(f" Batches per epoch: {n_batches_per_epoch:,}") -print(f" Samples per epoch: {n_batches_per_epoch * batch_size:,}") - -if n_batches_per_epoch > 20000: - print(f" ⚠️ WARNING: {n_batches_per_epoch:,} batches will be slow (~{n_batches_per_epoch/60:.0f} min/epoch at 1 batch/sec)") -else: - print(f" ✅ Reasonable number of batches per epoch") - -print() - -# ============================================================================ -# SUMMARY -# ============================================================================ -print("="*80) -print("SANITY CHECK SUMMARY") -print("="*80) - -checks = [ - ("Mapping file integrity", True), - ("Transitive closure indices", max_idx <= max_mapping_idx), - ("No self-loops", self_loops == 0), - ("Depth consistency", len(invalid_depths) == 0), - ("Positive depth diffs", negative_depths == 0), - ("Projection logic", True), # Tested above - ("Hyperbolic distance", True), # Tested above - ("Initialization radii", r_max_global < 1.0), - ("Sibling map", nodes_with_no_siblings < total_nodes_with_siblings * 0.5), - ("Regularizer targets", invalid_targets == 0), -] - -passed = sum(1 for _, status in checks if status) -total = len(checks) - -print(f"\nPassed {passed}/{total} checks") -print() - -for name, status in checks: - status_str = "✅ PASS" if status else "❌ FAIL" - print(f" {status_str}: {name}") - -print() -if passed == total: - print("✅ ALL CHECKS PASSED - Ready to train!") -else: - print(f"❌ {total - passed} CHECKS FAILED - Fix issues before training!") -print() diff --git a/scripts/build_clade_dataset.py b/scripts/build_clade_dataset.py deleted file mode 100644 index 98ca6b7..0000000 --- a/scripts/build_clade_dataset.py +++ /dev/null @@ -1,73 +0,0 @@ -#!/usr/bin/env python3 -"""CLI wrapper around the TaxoPy-backed clade dataset builder.""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -PROJECT_ROOT = Path(__file__).resolve().parents[1] -SRC_DIR = PROJECT_ROOT / "src" -if str(SRC_DIR) not in sys.path: - sys.path.insert(0, str(SRC_DIR)) - -from taxembed.builders import build_clade_dataset - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Build a taxonomy dataset for a specific clade.") - parser.add_argument("--root-taxid", type=int, required=True, help="Root TaxID for the clade.") - parser.add_argument( - "--dataset-name", - type=str, - default=None, - help="Optional dataset name override (defaults to sanitized root name).", - ) - parser.add_argument( - "--output-dir", - type=Path, - default=PROJECT_ROOT / "data" / "taxopy", - help="Directory where the dataset artifacts will be written.", - ) - parser.add_argument( - "--taxdump-dir", - type=Path, - default=PROJECT_ROOT / "data", - help="Directory containing nodes.dmp/names.dmp (downloaded automatically if missing).", - ) - parser.add_argument( - "--max-depth", - type=int, - default=None, - help="Optional depth limit (relative to the root taxon).", - ) - return parser.parse_args() - - -def main() -> None: - args = parse_args() - result = build_clade_dataset( - args.root_taxid, - dataset_name=args.dataset_name, - output_dir=args.output_dir, - taxdump_dir=args.taxdump_dir, - max_depth=args.max_depth, - ) - - print("\n✅ Finished building clade dataset") - print(f" Dataset: {result.dataset_name}") - print(f" Root TaxID: {result.root_taxid}") - print(f" Nodes: {result.node_count:,}") - print(f" Edges: {result.edge_count:,}") - print(f" Max depth observed: {result.max_depth}") - print(f" Transitive pairs: {result.pairs_count:,}") - print(f" Output directory: {result.output_dir}") - print("\n Files:") - for key, path in result.files.items(): - print(f" - {key}: {path}") - - -if __name__ == "__main__": - main() - diff --git a/scripts/cleanup_repo.sh b/scripts/cleanup_repo.sh deleted file mode 100755 index 886850f..0000000 --- a/scripts/cleanup_repo.sh +++ /dev/null @@ -1,87 +0,0 @@ -#!/bin/bash -# Clean up repository: remove checkpoints, logs, and temporary files - -set -e - -echo "===========================================" -echo "REPOSITORY CLEANUP" -echo "===========================================" -echo "" - -# Count files to be removed -checkpoint_count=$(find . -name "*.pth" -o -name "*.pth.*" | wc -l | tr -d ' ') -log_count=$(find . -maxdepth 1 -name "*.log" | wc -l | tr -d ' ') -png_count=$(find . -maxdepth 1 -name "*.png" | wc -l | tr -d ' ') - -echo "Files to be removed:" -echo " - Checkpoints (*.pth, *.pth.*): $checkpoint_count files" -echo " - Logs (*.log): $log_count files" -echo " - Visualizations (*.png): $png_count files" -echo "" - -# Ask for confirmation -read -p "Continue with cleanup? (y/N) " -n 1 -r -echo "" -if [[ ! $REPLY =~ ^[Yy]$ ]]; then - echo "Cleanup cancelled." - exit 0 -fi - -echo "" -echo "Removing checkpoint files..." -find . -name "*.pth" -delete -find . -name "*.pth.*" -delete -echo "✓ Removed $checkpoint_count checkpoint files" - -echo "" -echo "Removing log files..." -find . -maxdepth 1 -name "*.log" -delete -echo "✓ Removed $log_count log files" - -echo "" -echo "Removing visualization files..." -find . -maxdepth 1 -name "*.png" -delete -echo "✓ Removed $png_count PNG files" - -echo "" -echo "Removing redundant scripts..." -# Keep only the consolidated visualization script -rm -f visualize_primates.py -rm -f visualize_primates_proper.py -rm -f visualize_primates_small_only.py -rm -f visualize_by_taxonomy.py -rm -f visualize_trained_small_dataset.py -echo "✓ Removed redundant visualization scripts" - -echo "" -echo "Removing old shell scripts..." -rm -f train-mammals.sh -rm -f train-nouns.sh -rm -f train_taxonomy.sh -rm -f train_taxonomy_quick.sh -echo "✓ Removed old training scripts" - -echo "" -echo "Removing nohup.out..." -rm -f nohup.out -echo "✓ Removed nohup.out" - -echo "" -echo "===========================================" -echo "✅ CLEANUP COMPLETE" -echo "===========================================" -echo "" -echo "Repository is now clean and organized!" -echo "" -echo "Remaining structure:" -echo " scripts/ - Organized scripts" -echo " src/taxembed/ - Source code" -echo " tests/ - Unit tests" -echo " data/ - Data files (gitignored)" -echo " docs/ - Documentation" -echo "" -echo "To train a model:" -echo " python embed.py -dset data/taxonomy_edges_small.mapped.edgelist ..." -echo "" -echo "To visualize embeddings:" -echo " python scripts/visualize_embeddings.py --highlight primates" diff --git a/scripts/evaluate.py b/scripts/evaluate.py deleted file mode 100644 index 245dca2..0000000 --- a/scripts/evaluate.py +++ /dev/null @@ -1,16 +0,0 @@ -#!/usr/bin/env python3 -"""Evaluate trained embeddings. - -Computes reconstruction metrics and other evaluation measures. -""" - -import os -import sys - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) - -# Import and run the original script -from evaluate_full import main - -if __name__ == "__main__": - main() diff --git a/scripts/monitor.py b/scripts/monitor.py deleted file mode 100644 index 084158a..0000000 --- a/scripts/monitor.py +++ /dev/null @@ -1,16 +0,0 @@ -#!/usr/bin/env python3 -"""Monitor training progress in real-time. - -Displays clustering quality metrics during training. -""" - -import os -import sys - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) - -# Import and run the original script -from monitor_training import main - -if __name__ == "__main__": - main() diff --git a/scripts/prepare_data.py b/scripts/prepare_data.py deleted file mode 100644 index 27816cc..0000000 --- a/scripts/prepare_data.py +++ /dev/null @@ -1,16 +0,0 @@ -#!/usr/bin/env python3 -"""Prepare NCBI taxonomy data for embedding. - -Downloads and processes NCBI taxonomy data into edge list format. -""" - -import os -import sys - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) - -# Import and run the original script -from prepare_taxonomy_data import main - -if __name__ == "__main__": - main() diff --git a/scripts/regenerate_data.sh b/scripts/regenerate_data.sh deleted file mode 100755 index 34c0d3d..0000000 --- a/scripts/regenerate_data.sh +++ /dev/null @@ -1,66 +0,0 @@ -#!/bin/bash -# Regenerate clean training data from NCBI taxonomy - -set -e # Exit on error - -echo "=========================================" -echo "REGENERATING CLEAN TRAINING DATA" -echo "=========================================" -echo "" - -# Check if data directory exists -if [ ! -d "data" ]; then - echo "❌ data/ directory not found" - exit 1 -fi - -# Check if nodes.dmp exists -if [ ! -f "data/nodes.dmp" ]; then - echo "❌ data/nodes.dmp not found" - echo "Please download NCBI taxonomy first:" - echo " wget https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/new_taxdump/new_taxdump.tar.gz" - echo " tar -xzf new_taxdump.tar.gz -C data/" - exit 1 -fi - -echo "Step 1: Parsing NCBI taxonomy..." -python prepare_taxonomy_data.py -echo "" - -echo "Step 2: Creating small subset..." -# Create small subset (first 100k edges) -if [ -f "data/taxonomy_edges.edgelist" ]; then - head -n 100001 data/taxonomy_edges.edgelist > data/taxonomy_edges_small.edgelist - echo "✓ Created data/taxonomy_edges_small.edgelist" -else - echo "⚠️ data/taxonomy_edges.edgelist not found, skipping small subset" -fi -echo "" - -echo "Step 3: Remapping edges (full dataset)..." -python remap_edges.py data/taxonomy_edges.edgelist -echo "" - -echo "Step 4: Remapping edges (small dataset)..." -python remap_edges.py data/taxonomy_edges_small.edgelist -echo "" - -echo "Step 5: Validating data..." -python scripts/validate_data.py full -python scripts/validate_data.py small -echo "" - -echo "=========================================" -echo "✅ DATA REGENERATION COMPLETE" -echo "=========================================" -echo "" -echo "Files created:" -echo " - data/taxonomy_edges.csv" -echo " - data/taxonomy_edges.edgelist" -echo " - data/taxonomy_edges.mapped.edgelist" -echo " - data/taxonomy_edges.mapping.tsv" -echo " - data/taxonomy_edges_small.edgelist" -echo " - data/taxonomy_edges_small.mapped.edgelist" -echo " - data/taxonomy_edges_small.mapping.tsv" -echo "" -echo "Ready to train!" diff --git a/scripts/remap_data.py b/scripts/remap_data.py deleted file mode 100644 index cda840e..0000000 --- a/scripts/remap_data.py +++ /dev/null @@ -1,16 +0,0 @@ -#!/usr/bin/env python3 -"""Remap taxonomy IDs in edge list. - -Converts original taxonomy IDs to sequential indices for efficient training. -""" - -import os -import sys - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) - -# Import and run the original script -from remap_edges import main - -if __name__ == "__main__": - main() diff --git a/scripts/train.py b/scripts/train.py deleted file mode 100644 index fa4bbe6..0000000 --- a/scripts/train.py +++ /dev/null @@ -1,230 +0,0 @@ -#!/usr/bin/env python3 -"""Training script for Poincaré embeddings. - -This script trains hierarchical embeddings on graph data using Poincaré geometry. -""" - -import os -import sys - -# Suppress PyTorch verbose logging on startup -os.environ["TORCH_CPP_LOG_LEVEL"] = "ERROR" -os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" - -# Add parent directory to path for imports -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) - -import argparse -import json -import logging -import shutil -import torch as th -import torch.multiprocessing as mp -import numpy as np - -try: - from hype.adjacency_matrix_dataset import AdjacencyDataset -except Exception: - AdjacencyDataset = None - -from hype import MANIFOLDS, MODELS, build_model, train -from hype.checkpoint import LocalCheckpoint -from hype.graph import eval_reconstruction, load_adjacency_matrix, load_edge_list -from hype.graph_dataset import BatchedDataset -from hype.rsgd import RiemannianSGD - -# Optional import for hypernymy evaluation -try: - from hype.hypernymy_eval import main as hype_eval -except ImportError: - hype_eval = None - -th.manual_seed(42) -np.random.seed(42) - - -def reconstruction_eval(adj, opt, epoch, elapsed, loss, pth, best): - """Evaluate reconstruction metrics.""" - chkpnt = th.load(pth, map_location="cpu") - model = build_model(opt, chkpnt["embeddings"].size(0)) - model.load_state_dict(chkpnt["model"]) - - meanrank, maprank = eval_reconstruction(adj, model) - sqnorms = model.manifold.norm(model.lt) - return { - "epoch": epoch, - "elapsed": elapsed, - "loss": loss, - "mean_rank": meanrank.item(), - "map": maprank.item(), - "sqnorm_min": sqnorms.min().item(), - "sqnorm_max": sqnorms.max().item(), - "sqnorm_mean": sqnorms.mean().item(), - } - - -def main(): - """Main training function.""" - parser = argparse.ArgumentParser(description="Train Poincaré embeddings") - parser.add_argument( - "-dset", - "--dataset", - type=str, - required=True, - help="Path to dataset file (edgelist format)", - ) - parser.add_argument( - "-checkpoint", - "--checkpoint", - type=str, - required=True, - help="Path to save checkpoint", - ) - parser.add_argument( - "-dim", "--dim", type=int, default=10, help="Embedding dimension" - ) - parser.add_argument( - "-epochs", "--epochs", type=int, default=50, help="Number of epochs" - ) - parser.add_argument( - "-negs", - "--negs", - type=int, - default=50, - help="Number of negative samples", - ) - parser.add_argument( - "-burnin", - "--burnin", - type=int, - default=10, - help="Burn-in period", - ) - parser.add_argument( - "-batchsize", - "--batchsize", - type=int, - default=32, - help="Batch size", - ) - parser.add_argument( - "-model", - "--model", - type=str, - default="distance", - choices=MODELS.keys(), - help="Model type", - ) - parser.add_argument( - "-manifold", - "--manifold", - type=str, - default="poincare", - choices=MANIFOLDS.keys(), - help="Manifold type", - ) - parser.add_argument( - "-lr", "--lr", type=float, default=0.1, help="Learning rate" - ) - parser.add_argument( - "-gpu", "--gpu", type=int, default=-1, help="GPU ID (-1 for CPU)" - ) - parser.add_argument( - "-ndproc", - "--ndproc", - type=int, - default=1, - help="Number of data loading processes", - ) - parser.add_argument( - "-train_threads", - "--train_threads", - type=int, - default=1, - help="Number of training threads", - ) - parser.add_argument( - "-eval_each", - "--eval_each", - type=int, - default=999999, - help="Evaluate every N epochs", - ) - parser.add_argument( - "-fresh", - "--fresh", - action="store_true", - help="Start fresh training", - ) - - args = parser.parse_args() - - # Setup logging - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(levelname)s - %(message)s", - ) - logger = logging.getLogger(__name__) - - logger.info(f"Loading dataset from {args.dataset}") - - # Load data - if args.dataset.endswith(".csv"): - adj = load_adjacency_matrix(args.dataset, "csv") - else: - adj = load_edge_list(args.dataset) - - logger.info(f"Dataset loaded: {adj.shape[0]} nodes, {adj.nnz} edges") - - # Build model - opt = { - "model": args.model, - "manifold": args.manifold, - "dim": args.dim, - "epochs": args.epochs, - "negs": args.negs, - "burnin": args.burnin, - "batchsize": args.batchsize, - "lr": args.lr, - "gpu": args.gpu, - "ndproc": args.ndproc, - "train_threads": args.train_threads, - "eval_each": args.eval_each, - } - - model = build_model(opt, adj.shape[0]) - logger.info(f"Model built: {args.model} on {args.manifold}") - - # Setup optimizer - optimizer = RiemannianSGD(model.parameters(), lr=args.lr) - - # Setup checkpoint - checkpoint = LocalCheckpoint(args.checkpoint, include_in_all=["model"]) - - # Load checkpoint if exists and not fresh - if not args.fresh and os.path.exists(args.checkpoint): - logger.info(f"Loading checkpoint from {args.checkpoint}") - chkpnt = th.load(args.checkpoint, map_location="cpu") - model.load_state_dict(chkpnt["model"]) - optimizer.load_state_dict(chkpnt["optimizer"]) - - # Setup dataset - dataset = BatchedDataset(adj, opt["batchsize"], opt["negs"]) - - # Train - logger.info("Starting training...") - train( - model, - optimizer, - dataset, - opt, - checkpoint, - reconstruction_eval if args.eval_each < 999999 else None, - adj if args.eval_each < 999999 else None, - ) - - logger.info(f"Training complete. Model saved to {args.checkpoint}") - - -if __name__ == "__main__": - main() diff --git a/scripts/validate_data.py b/scripts/validate_data.py deleted file mode 100644 index 03d55bc..0000000 --- a/scripts/validate_data.py +++ /dev/null @@ -1,239 +0,0 @@ -#!/usr/bin/env python3 -""" -Validate data files for Poincaré embedding training. - -Checks: -1. Edgelist files have no headers -2. All values are numeric -3. Mapping files are consistent -4. Node indices are sequential -""" - -from __future__ import annotations - -import sys -from pathlib import Path - -import pandas as pd - -ROOT_DIR = Path(__file__).resolve().parents[1] -SRC_DIR = ROOT_DIR / "src" -if str(SRC_DIR) not in sys.path: - sys.path.insert(0, str(SRC_DIR)) - -from taxembed.utils.data_validation import coverage_from_indices, load_mapping - - -def validate_edgelist(filepath): - """Validate an edgelist file.""" - print(f"\n📋 Validating edgelist: {filepath}") - - issues = [] - edges = [] - nodes = set() - - with open(filepath, 'r') as f: - for ln, line in enumerate(f, 1): - line = line.strip() - if not line: - continue - - parts = line.split() - if len(parts) != 2: - issues.append(f"Line {ln}: Expected 2 values, got {len(parts)}") - continue - - # Check if numeric - try: - u, v = int(parts[0]), int(parts[1]) - edges.append((u, v)) - nodes.add(u) - nodes.add(v) - except ValueError: - issues.append(f"Line {ln}: Non-numeric values: {parts}") - - # Check for headers - if issues and issues[0].startswith("Line 1"): - print(" ⚠️ Possible header line detected") - - # Statistics - print(f" ✓ Total edges: {len(edges):,}") - print(f" ✓ Unique nodes: {len(nodes):,}") - - if len(nodes) > 0: - min_node = min(nodes) - max_node = max(nodes) - print(f" ✓ Node range: {min_node} to {max_node}") - - # Check if sequential - if max_node - min_node + 1 == len(nodes): - print(f" ✓ Nodes are sequential") - else: - expected = max_node - min_node + 1 - print(f" ⚠️ Nodes are NOT sequential (expected {expected}, got {len(nodes)})") - - # Report issues - if issues: - print(f"\n ❌ Found {len(issues)} issues:") - for issue in issues[:10]: # Show first 10 - print(f" {issue}") - if len(issues) > 10: - print(f" ... and {len(issues) - 10} more") - return False - else: - print(f" ✅ No issues found") - return True - - -def validate_mapping(filepath): - """Validate a mapping file.""" - print(f"\n📋 Validating mapping: {filepath}") - - try: - df = load_mapping(Path(filepath)) - except Exception as e: - print(f" ❌ Failed to read: {e}") - return False - - print(f" ✓ Total mappings: {len(df):,}") - - invalid_mask = ~df["taxid"].str.isnumeric() - non_numeric = df[invalid_mask] - - if not non_numeric.empty: - print(f" ❌ Found {len(non_numeric)} non-numeric TaxIDs:") - for idx, row in non_numeric.head(10).iterrows(): - print(f" Row {idx}: taxid='{row['taxid']}' is not numeric") - return False - - # Check if indices are sequential - indices = sorted(df['idx'].values) - if indices == list(range(len(indices))): - print(f" ✓ Indices are sequential (0 to {len(indices)-1})") - else: - print(f" ⚠️ Indices are NOT sequential") - print(f" Expected: 0 to {len(indices)-1}") - print(f" Got: {indices[0]} to {indices[-1]}") - - # Check for duplicates - dup_taxids = df[df.duplicated('taxid', keep=False)] - dup_indices = df[df.duplicated('idx', keep=False)] - - if len(dup_taxids) > 0: - print(f" ❌ Found {len(dup_taxids)} duplicate TaxIDs") - return False - if len(dup_indices) > 0: - print(f" ❌ Found {len(dup_indices)} duplicate indices") - return False - - print(f" ✅ No issues found") - return True - - -def validate_consistency(edgelist_file, mapping_file): - """Validate consistency between edgelist and mapping.""" - print(f"\n📋 Validating consistency...") - - # Load edgelist nodes - nodes = set() - with open(edgelist_file, 'r') as f: - for line in f: - line = line.strip() - if not line: - continue - parts = line.split() - if len(parts) == 2: - try: - nodes.add(int(parts[0])) - nodes.add(int(parts[1])) - except ValueError: - pass - - # Load mapping - df = load_mapping(Path(mapping_file)) - - # Filter out non-numeric taxids - numeric_df = df[df['taxid'].str.isnumeric()] - mapped_indices = set(numeric_df['idx'].values) - - print(f" Edgelist nodes: {len(nodes):,}") - print(f" Mapping indices: {len(mapped_indices):,}") - - # Check if all edgelist nodes are in mapping - unmapped = nodes - mapped_indices - if unmapped: - print(f" ⚠️ {len(unmapped)} nodes in edgelist not in mapping") - print(f" Examples: {sorted(list(unmapped))[:10]}") - else: - print(f" ✓ All edgelist nodes are in mapping") - - report = coverage_from_indices(numeric_df, nodes) - if report.is_perfect: - print(f" ✓ All mapping indices are used in edgelist") - else: - print( - f" ⚠️ Coverage gap: {report.missing_count} indices missing " - f"({report.coverage_ratio * 100:.1f}% covered)" - ) - sample = sorted(report.missing_indices)[:10] - if sample: - print(f" Examples: {sample}") - - return len(unmapped) == 0 - - -def main(): - if len(sys.argv) < 2: - print("Usage: python validate_data.py ") - print("Example: python validate_data.py small") - print(" python validate_data.py full") - sys.exit(1) - - dataset = sys.argv[1] - data_dir = Path(__file__).parent.parent / 'data' - - print(f"{'='*60}") - print(f"DATA VALIDATION - {dataset.upper()} DATASET") - print(f"{'='*60}") - - if dataset == 'full': - edgelist = data_dir / 'taxonomy_edges.mapped.edgelist' - mapping = data_dir / 'taxonomy_edges.mapping.tsv' - elif dataset == 'small': - edgelist = data_dir / 'taxonomy_edges_small.mapped.edgelist' - mapping = data_dir / 'taxonomy_edges_small.mapping.tsv' - else: - print(f"Unknown dataset: {dataset}") - sys.exit(1) - - # Validate files - results = [] - - if edgelist.exists(): - results.append(validate_edgelist(edgelist)) - else: - print(f"\n❌ File not found: {edgelist}") - results.append(False) - - if mapping.exists(): - results.append(validate_mapping(mapping)) - else: - print(f"\n❌ File not found: {mapping}") - results.append(False) - - if edgelist.exists() and mapping.exists(): - results.append(validate_consistency(edgelist, mapping)) - - # Summary - print(f"\n{'='*60}") - if all(results): - print("✅ ALL CHECKS PASSED") - else: - print("❌ SOME CHECKS FAILED") - print(f"{'='*60}\n") - - return 0 if all(results) else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/scripts/visualize.py b/scripts/visualize.py deleted file mode 100644 index 2382495..0000000 --- a/scripts/visualize.py +++ /dev/null @@ -1,16 +0,0 @@ -#!/usr/bin/env python3 -"""Visualize embeddings using UMAP. - -Creates 2D projections of the learned embeddings for visualization. -""" - -import os -import sys - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) - -# Import and run the original script -from evaluate_and_visualize import main - -if __name__ == "__main__": - main() diff --git a/scripts/visualize_embeddings.py b/scripts/visualize_embeddings.py deleted file mode 100644 index abda8ff..0000000 --- a/scripts/visualize_embeddings.py +++ /dev/null @@ -1,346 +0,0 @@ -#!/usr/bin/env python3 -""" -Universal visualization tool for Poincaré embeddings. - -Features: -- Works with any checkpoint -- Can highlight specific taxonomic groups (primates, mammals, bacteria, etc.) -- Supports different sampling strategies -- Generates UMAP projections -- Nearest neighbor analysis - -Usage: - python scripts/visualize_embeddings.py [options] - -Examples: - # Basic visualization - python scripts/visualize_embeddings.py model.pth - - # Highlight primates - python scripts/visualize_embeddings.py model.pth --highlight primates - - # Highlight mammals with custom sample size - python scripts/visualize_embeddings.py model.pth --highlight mammals --sample 50000 - - # Only visualize specific group - python scripts/visualize_embeddings.py model.pth --only primates -""" - -import argparse -import sys -from pathlib import Path - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import torch -from umap import UMAP - - -def load_embeddings(ckpt_path): - """Load embeddings from checkpoint.""" - print(f"Loading embeddings from {ckpt_path}...") - ckpt = torch.load(ckpt_path, map_location="cpu") - - if "state_dict" in ckpt: - sd = ckpt["state_dict"] - emb = sd["lt.weight"].detach().cpu().numpy() - elif "embeddings" in ckpt: - emb = ckpt["embeddings"].cpu().numpy() - else: - raise ValueError("Cannot find embeddings in checkpoint") - - print(f" ✓ Shape: {emb.shape}") - return emb - - -def load_mapping(map_path): - """Load TaxID to index mapping.""" - if not Path(map_path).exists(): - print(f"⚠️ Mapping file not found: {map_path}") - return None, None - - print(f"Loading mapping from {map_path}...") - df = pd.read_csv(map_path, sep="\t", dtype={"taxid": str, "idx": int}) - - # Filter out non-numeric taxids - numeric_df = df[df["taxid"].str.isnumeric()] - tax2idx = dict(zip(numeric_df["taxid"], numeric_df["idx"])) - idx2tax = dict(zip(numeric_df["idx"], numeric_df["taxid"])) - - print(f" ✓ Loaded {len(tax2idx):,} mappings") - return tax2idx, idx2tax - - -def load_taxonomy_tree(valid_taxids=None): - """Load NCBI taxonomy tree structure, optionally filtered to valid TaxIDs. - - Args: - valid_taxids: Set of TaxIDs to include. If None, loads all. - """ - try: - # Load names - names = {} - with open("data/names.dmp", "r") as f: - for line in f: - parts = [p.strip() for p in line.split("|")] - if len(parts) >= 4 and parts[3] == "scientific name": - taxid = int(parts[0]) - # Only load if in valid set or loading all - if valid_taxids is None or taxid in valid_taxids: - names[taxid] = parts[1] - - # Load nodes (parent relationships) - nodes = {} - with open("data/nodes.dmp", "r") as f: - for line in f: - parts = [p.strip() for p in line.split("|")] - if len(parts) >= 5: - taxid = int(parts[0]) - parent = int(parts[1]) - rank = parts[2] - # Only load if in valid set or loading all - if valid_taxids is None or taxid in valid_taxids: - nodes[taxid] = {"parent": parent, "rank": rank, "name": names.get(taxid, "")} - - if valid_taxids: - print(f" ✓ Loaded {len(nodes):,} taxonomy nodes (filtered to dataset)") - else: - print(f" ✓ Loaded {len(nodes):,} taxonomy nodes") - return nodes - - except FileNotFoundError: - print(" ⚠️ Taxonomy files not found (data/nodes.dmp, data/names.dmp)") - return {} - - -def find_taxonomic_group(nodes, group_name): - """Find all TaxIDs in a taxonomic group.""" - # Common groups and their NCBI TaxIDs - group_roots = { - "primates": 9443, - "mammals": 40674, - "mammalia": 40674, - "vertebrates": 7742, - "bacteria": 2, - "archaea": 2157, - "fungi": 4751, - "plants": 33090, - "insects": 50557, - "rodents": 9989, - } - - group_name = group_name.lower() - if group_name not in group_roots: - print(f" ⚠️ Unknown group: {group_name}") - print(f" Available groups: {', '.join(group_roots.keys())}") - return set() - - root_taxid = group_roots[group_name] - - # Find all descendants - print(f"Finding {group_name} (root TaxID: {root_taxid})...") - group_taxids = set() - - def find_descendants(taxid): - group_taxids.add(taxid) - for child_id, child_info in nodes.items(): - if child_info["parent"] == taxid: - find_descendants(child_id) - - find_descendants(root_taxid) - print(f" ✓ Found {len(group_taxids):,} {group_name}") - return group_taxids - - -def nearest_neighbors(emb, idx, k=10): - """Find k nearest neighbors.""" - x = emb[idx] - d = np.linalg.norm(emb - x, axis=1) - nbrs = np.argsort(d)[:k+1] - nbrs = [j for j in nbrs if j != idx][:k] - return nbrs, d[nbrs] - - -def visualize_embeddings(emb, indices, colors, labels, title, output_file, sample_size=None): - """Create UMAP visualization.""" - # Sample if needed - if sample_size and len(indices) > sample_size: - print(f"Sampling {sample_size:,} points from {len(indices):,} total...") - sample_idx = np.random.choice(len(indices), sample_size, replace=False) - indices = [indices[i] for i in sample_idx] - colors = [colors[i] for i in sample_idx] - - # Extract embeddings - sample_emb = emb[indices] - - # Run UMAP - print(f"Running UMAP on {len(indices):,} points...") - umap_model = UMAP(n_components=2, random_state=42, n_neighbors=15, min_dist=0.1) - projection = umap_model.fit_transform(sample_emb) - - # Plot - print("Creating visualization...") - fig, ax = plt.subplots(figsize=(16, 12)) - - # Plot points - for label, color in labels.items(): - mask = np.array([c == color for c in colors]) - if mask.sum() > 0: - ax.scatter( - projection[mask, 0], - projection[mask, 1], - c=color, - s=30 if label == "Highlighted" else 20, - alpha=0.7 if label == "Highlighted" else 0.3, - label=f"{label} (n={mask.sum():,})", - edgecolors="darkred" if label == "Highlighted" else "none", - linewidth=0.5 if label == "Highlighted" else 0, - ) - - ax.set_xlabel("UMAP 1", fontsize=14) - ax.set_ylabel("UMAP 2", fontsize=14) - ax.set_title(title, fontsize=16, fontweight="bold") - ax.legend(loc="best", fontsize=11, framealpha=0.9) - ax.grid(True, alpha=0.3) - - plt.tight_layout() - plt.savefig(output_file, dpi=200, bbox_inches="tight") - print(f"✓ Saved: {output_file}") - - -def main(): - parser = argparse.ArgumentParser( - description="Visualize Poincaré embeddings", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=__doc__, - ) - parser.add_argument("checkpoint", help="Path to checkpoint file") - parser.add_argument( - "--mapping", - help="Path to mapping file (auto-detected if not specified)", - ) - parser.add_argument( - "--highlight", - help="Taxonomic group to highlight (primates, mammals, bacteria, etc.)", - ) - parser.add_argument( - "--only", - help="Only show specific taxonomic group", - ) - parser.add_argument( - "--sample", - type=int, - default=25000, - help="Number of points to sample (default: 25000)", - ) - parser.add_argument( - "--output", - help="Output filename (default: auto-generated)", - ) - parser.add_argument( - "--nearest", - type=int, - default=5, - help="Show N nearest neighbors for key organisms (default: 5)", - ) - - args = parser.parse_args() - - # Load embeddings - emb = load_embeddings(args.checkpoint) - n_nodes = emb.shape[0] - - # Auto-detect mapping file - if not args.mapping: - checkpoint_path = Path(args.checkpoint) - - # Try to infer from checkpoint name - if "small" in checkpoint_path.name: - args.mapping = "data/taxonomy_edges_small.mapping.tsv" - else: - args.mapping = "data/taxonomy_edges.mapping.tsv" - - # Load mapping - tax2idx, idx2tax = load_mapping(args.mapping) - - # Get valid TaxIDs from the mapping (only organisms in training data) - valid_taxids = set(int(taxid) for taxid in idx2tax.values()) - print(f"Dataset contains {len(valid_taxids):,} unique organisms") - - # Load taxonomy if highlighting (filtered to dataset organisms only) - nodes = {} - if args.highlight or args.only: - print("Loading taxonomy tree...") - nodes = load_taxonomy_tree(valid_taxids=valid_taxids) - - # Determine which indices to visualize - if args.only: - # Only show specific group - group_taxids = find_taxonomic_group(nodes, args.only) - indices = [idx for idx, tax in idx2tax.items() if int(tax) in group_taxids] - colors = ["red"] * len(indices) - labels = {"Highlighted": "red"} - title = f"{args.only.capitalize()} - UMAP Projection\n({len(indices):,} organisms)" - elif args.highlight: - # Show all, highlight specific group - group_taxids = find_taxonomic_group(nodes, args.highlight) - indices = list(range(n_nodes)) - colors = ["red" if idx2tax.get(idx) and int(idx2tax[idx]) in group_taxids else "lightgray" for idx in indices] - n_highlighted = colors.count("red") - labels = {"Highlighted": "red", "Other": "lightgray"} - title = f"Embedding Visualization - {args.highlight.capitalize()} Highlighted\n({n_highlighted:,} highlighted, {n_nodes-n_highlighted:,} other)" - else: - # Show all - indices = list(range(min(n_nodes, args.sample))) - colors = ["steelblue"] * len(indices) - labels = {"All": "steelblue"} - title = f"Embedding Visualization\n({len(indices):,} organisms)" - - # Generate output filename - if not args.output: - checkpoint_name = Path(args.checkpoint).stem - if args.only: - args.output = f"umap_{checkpoint_name}_{args.only}_only.png" - elif args.highlight: - args.output = f"umap_{checkpoint_name}_{args.highlight}_highlighted.png" - else: - args.output = f"umap_{checkpoint_name}.png" - - # Visualize - visualize_embeddings( - emb, indices, colors, labels, title, args.output, args.sample if not args.only else None - ) - - # Nearest neighbors analysis - if tax2idx and args.nearest > 0: - print(f"\n{'='*70}") - print("NEAREST NEIGHBORS ANALYSIS") - print(f"{'='*70}\n") - - key_organisms = { - "9606": "Homo sapiens (Human)", - "10090": "Mus musculus (Mouse)", - "6239": "Caenorhabditis elegans", - "7227": "Drosophila melanogaster", - "562": "Escherichia coli", - } - - for taxid, name in key_organisms.items(): - if taxid in tax2idx: - idx = tax2idx[taxid] - nbrs, dists = nearest_neighbors(emb, idx, args.nearest) - print(f"{name}:") - for i, (nbr, dist) in enumerate(zip(nbrs, dists), 1): - nbr_tax = idx2tax.get(nbr, "Unknown") - print(f" {i}. TaxID {nbr_tax} (distance: {dist:.6f})") - print() - - print(f"\n{'='*70}") - print("✅ VISUALIZATION COMPLETE") - print(f"{'='*70}") - print(f"Output: {args.output}") - - -if __name__ == "__main__": - main() diff --git a/src/taxembed/__init__.py b/src/taxembed/__init__.py index 95be8e4..e8d2156 100644 --- a/src/taxembed/__init__.py +++ b/src/taxembed/__init__.py @@ -7,6 +7,8 @@ __author__ = "Your Name" __email__ = "your.email@example.com" -from . import manifolds # noqa: F401 -from . import models # noqa: F401 -from . import utils # noqa: F401 +from . import ( + manifolds, # noqa: F401 + models, # noqa: F401 + utils, # noqa: F401 +) diff --git a/src/taxembed/analysis/__init__.py b/src/taxembed/analysis/__init__.py new file mode 100644 index 0000000..b1e45df --- /dev/null +++ b/src/taxembed/analysis/__init__.py @@ -0,0 +1,7 @@ +"""Analysis tools for embeddings quality and hierarchy.""" + +from .hierarchy import main as analyze_hierarchy + +__all__ = [ + "analyze_hierarchy", +] diff --git a/analyze_hierarchy_hyperbolic.py b/src/taxembed/analysis/hierarchy.py similarity index 76% rename from analyze_hierarchy_hyperbolic.py rename to src/taxembed/analysis/hierarchy.py index f49f743..990efe2 100644 --- a/analyze_hierarchy_hyperbolic.py +++ b/src/taxembed/analysis/hierarchy.py @@ -4,40 +4,41 @@ This is the correct metric for Poincaré embeddings, not Euclidean! """ -import torch -import numpy as np -import pandas as pd from collections import defaultdict + import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import torch def poincare_distance(u, v, eps=1e-5): """ Compute Poincaré distance between points u and v. - + d(u,v) = arcosh(1 + 2 * ||u-v||^2 / ((1-||u||^2)(1-||v||^2))) - + This is the TRUE hyperbolic distance in the Poincaré ball model. """ # Compute squared norms u_norm_sq = np.sum(u**2) v_norm_sq = np.sum(v**2) - + # Clamp to stay inside the ball (norm < 1) u_norm_sq = min(u_norm_sq, 1 - eps) v_norm_sq = min(v_norm_sq, 1 - eps) - + # Compute squared distance - diff_norm_sq = np.sum((u - v)**2) - + diff_norm_sq = np.sum((u - v) ** 2) + # Poincaré distance formula numerator = 2 * diff_norm_sq denominator = (1 - u_norm_sq) * (1 - v_norm_sq) - + # arcosh(1 + x) with numerical stability x = numerator / (denominator + eps) dist = np.arccosh(1 + x + eps) - + return dist @@ -45,13 +46,13 @@ def poincare_distance_matrix(embeddings, indices): """Compute pairwise Poincaré distances for a subset of embeddings.""" n = len(indices) distances = np.zeros((n, n)) - + for i in range(n): - for j in range(i+1, n): + for j in range(i + 1, n): dist = poincare_distance(embeddings[indices[i]], embeddings[indices[j]]) distances[i, j] = dist distances[j, i] = dist - + return distances @@ -59,7 +60,7 @@ def load_embeddings(ckpt_path): """Load embeddings from checkpoint.""" print(f"Loading embeddings from {ckpt_path}...") ckpt = torch.load(ckpt_path, map_location="cpu") - + if "state_dict" in ckpt: sd = ckpt["state_dict"] emb = sd["lt.weight"].detach().cpu().numpy() @@ -67,15 +68,15 @@ def load_embeddings(ckpt_path): emb = ckpt["embeddings"].cpu().numpy() else: raise ValueError("Cannot find embeddings in checkpoint") - + print(f" ✓ Shape: {emb.shape}") - + # Check norms norms = np.linalg.norm(emb, axis=1) print(f" Norms: min={norms.min():.4f}, mean={norms.mean():.4f}, max={norms.max():.4f}") if norms.max() >= 1.0: - print(f" ⚠️ WARNING: Some embeddings are outside the Poincaré ball (norm >= 1.0)!") - + print(" ⚠️ WARNING: Some embeddings are outside the Poincaré ball (norm >= 1.0)!") + return emb @@ -84,7 +85,7 @@ def load_mapping(mapping_file): print(f"Loading mapping from {mapping_file}...") df = pd.read_csv(mapping_file, sep="\t", header=None, names=["idx", "taxid"]) numeric_df = df[df["taxid"].str.isnumeric()] - idx2tax = dict(zip(numeric_df["idx"], numeric_df["taxid"])) + idx2tax = dict(zip(numeric_df["idx"], numeric_df["taxid"], strict=False)) print(f" ✓ Loaded {len(idx2tax):,} mappings") return idx2tax @@ -92,20 +93,20 @@ def load_mapping(mapping_file): def load_taxonomy_with_depth(valid_taxids): """Load taxonomy and compute depth for each node.""" print("Loading taxonomy tree with depths...") - + # Load names names = {} - with open("data/names.dmp", "r") as f: + with open("data/names.dmp") as f: for line in f: parts = [p.strip() for p in line.split("|")] if len(parts) >= 4 and parts[3] == "scientific name": taxid = int(parts[0]) if taxid in valid_taxids: names[taxid] = parts[1] - + # Load nodes with ranks taxonomy = {} - with open("data/nodes.dmp", "r") as f: + with open("data/nodes.dmp") as f: for line in f: parts = [p.strip() for p in line.split("|")] if len(parts) >= 5: @@ -117,37 +118,37 @@ def load_taxonomy_with_depth(valid_taxids): "parent": parent, "rank": rank, "name": names.get(taxid, f"TaxID_{taxid}"), - "depth": None # Will compute + "depth": None, # Will compute } - + # Compute depth for each node def get_depth(taxid, visited=None): if visited is None: visited = set() - + if taxid not in taxonomy: return 0 - + if taxonomy[taxid]["depth"] is not None: return taxonomy[taxid]["depth"] - + if taxid in visited: # Cycle detection return 0 - + visited.add(taxid) parent = taxonomy[taxid]["parent"] - + if parent == taxid: # Root taxonomy[taxid]["depth"] = 0 else: taxonomy[taxid]["depth"] = get_depth(parent, visited) + 1 - + return taxonomy[taxid]["depth"] - + for taxid in taxonomy: if taxonomy[taxid]["depth"] is None: get_depth(taxid) - + print(f" ✓ Loaded {len(taxonomy):,} taxonomy nodes with depths") return taxonomy @@ -156,63 +157,64 @@ def get_ancestor_at_rank(taxid, taxonomy, target_rank): """Find the ancestor of a taxid at a specific rank.""" visited = set() current = taxid - + while current in taxonomy and current not in visited: visited.add(current) node = taxonomy[current] - + if node["rank"] == target_rank: return current - + parent = node["parent"] if parent == current: # Root break current = parent - + return None def analyze_depth_vs_norm(emb, idx2tax, taxonomy): """Check if depth correlates with radial position (CRITICAL check).""" - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print("RADIAL MONOTONICITY CHECK") - print(f"{'='*80}\n") - + print(f"{'=' * 80}\n") + max_idx = emb.shape[0] - 1 - + depths = [] norms = [] - + for idx_str, taxid_str in idx2tax.items(): idx = int(idx_str) if idx > max_idx: continue - + taxid = int(taxid_str) if taxid not in taxonomy: continue - + depth = taxonomy[taxid]["depth"] if depth is None: continue - + norm = np.linalg.norm(emb[idx]) - + depths.append(depth) norms.append(norm) - + depths = np.array(depths) norms = np.array(norms) - + # Compute correlation from scipy.stats import pearsonr, spearmanr + pearson_r, pearson_p = pearsonr(depths, norms) spearman_r, spearman_p = spearmanr(depths, norms) - - print(f"Depth vs. Norm correlation:") + + print("Depth vs. Norm correlation:") print(f" Pearson: r = {pearson_r:+.4f} (p = {pearson_p:.2e})") print(f" Spearman: ρ = {spearman_r:+.4f} (p = {spearman_p:.2e})") - + if pearson_r > 0.5: assessment = "✅ GOOD - Deeper nodes are near boundary" elif pearson_r > 0.3: @@ -221,98 +223,108 @@ def analyze_depth_vs_norm(emb, idx2tax, taxonomy): assessment = "❌ POOR - Very weak depth structure" else: assessment = "❌ BROKEN - Negative correlation! Hierarchy is inverted" - + print(f"\nAssessment: {assessment}") - + # Plot fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6)) - + # Scatter plot with density - ax1.hexbin(depths, norms, gridsize=50, cmap='YlOrRd', mincnt=1) - ax1.set_xlabel('Taxonomic Depth', fontsize=12) - ax1.set_ylabel('Embedding Norm (Distance from Origin)', fontsize=12) - ax1.set_title(f'Depth vs. Radial Position\nPearson r = {pearson_r:.3f}', - fontsize=14, fontweight='bold') + ax1.hexbin(depths, norms, gridsize=50, cmap="YlOrRd", mincnt=1) + ax1.set_xlabel("Taxonomic Depth", fontsize=12) + ax1.set_ylabel("Embedding Norm (Distance from Origin)", fontsize=12) + ax1.set_title( + f"Depth vs. Radial Position\nPearson r = {pearson_r:.3f}", fontsize=14, fontweight="bold" + ) ax1.grid(True, alpha=0.3) - + # Add trend line z = np.polyfit(depths, norms, 1) p = np.poly1d(z) - ax1.plot(depths, p(depths), "r--", linewidth=2, alpha=0.8, - label=f'Trend: y = {z[0]:.4f}x + {z[1]:.4f}') + ax1.plot( + depths, + p(depths), + "r--", + linewidth=2, + alpha=0.8, + label=f"Trend: y = {z[0]:.4f}x + {z[1]:.4f}", + ) ax1.legend() - + # Box plot by depth bins depth_bins = np.percentile(depths, [0, 25, 50, 75, 100]) - depth_labels = [f'{int(depth_bins[i])}-{int(depth_bins[i+1])}' - for i in range(len(depth_bins)-1)] + depth_labels = [ + f"{int(depth_bins[i])}-{int(depth_bins[i + 1])}" for i in range(len(depth_bins) - 1) + ] depth_binned = np.digitize(depths, depth_bins[1:-1]) - - data_by_bin = [norms[depth_binned == i] for i in range(len(depth_bins)-1)] + + data_by_bin = [norms[depth_binned == i] for i in range(len(depth_bins) - 1)] bp = ax2.boxplot(data_by_bin, labels=depth_labels, patch_artist=True) - for patch in bp['boxes']: - patch.set_facecolor('lightblue') - ax2.set_xlabel('Depth Quartile', fontsize=12) - ax2.set_ylabel('Embedding Norm', fontsize=12) - ax2.set_title('Norm Distribution by Depth', fontsize=14, fontweight='bold') - ax2.grid(True, alpha=0.3, axis='y') - + for patch in bp["boxes"]: + patch.set_facecolor("lightblue") + ax2.set_xlabel("Depth Quartile", fontsize=12) + ax2.set_ylabel("Embedding Norm", fontsize=12) + ax2.set_title("Norm Distribution by Depth", fontsize=14, fontweight="bold") + ax2.grid(True, alpha=0.3, axis="y") + plt.tight_layout() - plt.savefig('depth_vs_norm_analysis.png', dpi=150, bbox_inches='tight') - print(f"\nSaved plot: depth_vs_norm_analysis.png") + plt.savefig("depth_vs_norm_analysis.png", dpi=150, bbox_inches="tight") + print("\nSaved plot: depth_vs_norm_analysis.png") plt.close() - + return pearson_r, spearman_r def analyze_hierarchical_clustering_hyperbolic(emb, idx2tax, taxonomy, rank="phylum"): """Analyze hierarchical clustering using HYPERBOLIC distance.""" - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print(f"HIERARCHICAL CLUSTERING - {rank.upper()} (HYPERBOLIC DISTANCE)") - print(f"{'='*80}\n") - + print(f"{'=' * 80}\n") + max_idx = emb.shape[0] - 1 - + # Map organisms to groups organism_to_group = {} group_names = {} - + for idx_str, taxid_str in idx2tax.items(): idx = int(idx_str) if idx > max_idx: continue - + taxid = int(taxid_str) ancestor = get_ancestor_at_rank(taxid, taxonomy, rank) if ancestor: organism_to_group[idx] = ancestor if ancestor not in group_names and ancestor in taxonomy: group_names[ancestor] = taxonomy[ancestor]["name"] - + # Group by ancestor groups = defaultdict(list) for idx, group_id in organism_to_group.items(): groups[group_id].append(idx) - + # Filter large groups min_size = 10 large_groups = {gid: indices for gid, indices in groups.items() if len(indices) >= min_size} - + print(f"Found {len(set(organism_to_group.values()))} distinct {rank}s") print(f"{rank.capitalize()}s with ≥{min_size} organisms: {len(large_groups)}") - + # Show top groups - group_sizes = [(gid, len(indices), group_names.get(gid, f"TaxID_{gid}")) - for gid, indices in large_groups.items()] + group_sizes = [ + (gid, len(indices), group_names.get(gid, f"TaxID_{gid}")) + for gid, indices in large_groups.items() + ] group_sizes.sort(key=lambda x: -x[1]) - + print(f"\nTop 10 {rank}s by organism count:") - for i, (gid, size, name) in enumerate(group_sizes[:10], 1): + for i, (_gid, size, name) in enumerate(group_sizes[:10], 1): print(f" {i:2d}. {name:40s}: {size:6,} organisms") - + # Compute hyperbolic distances (sample for efficiency) - print(f"\nComputing pairwise HYPERBOLIC distances...") - + print("\nComputing pairwise HYPERBOLIC distances...") + max_per_group = 100 # Limit for efficiency sampled_groups = {} for gid, indices in large_groups.items(): @@ -320,53 +332,53 @@ def analyze_hierarchical_clustering_hyperbolic(emb, idx2tax, taxonomy, rank="phy sampled_groups[gid] = list(np.random.choice(indices, max_per_group, replace=False)) else: sampled_groups[gid] = list(indices) - + intra_distances = [] inter_distances = [] - + group_list = list(sampled_groups.items()) - - for i, (gid1, indices1) in enumerate(group_list): + + for i, (_gid1, indices1) in enumerate(group_list): # Intra-group distances if len(indices1) >= 2: for ii in range(len(indices1)): - for jj in range(ii+1, min(ii+20, len(indices1))): # Limit pairs + for jj in range(ii + 1, min(ii + 20, len(indices1))): # Limit pairs dist = poincare_distance(emb[indices1[ii]], emb[indices1[jj]]) intra_distances.append(dist) - + # Inter-group distances (sample) - for j in range(i+1, min(i+10, len(group_list))): # Limit group pairs + for j in range(i + 1, min(i + 10, len(group_list))): # Limit group pairs gid2, indices2 = group_list[j] - + # Sample pairs for _ in range(min(100, len(indices1) * len(indices2))): idx1 = np.random.choice(indices1) idx2 = np.random.choice(indices2) dist = poincare_distance(emb[idx1], emb[idx2]) inter_distances.append(dist) - + intra_distances = np.array(intra_distances) inter_distances = np.array(inter_distances) - - print(f"\n{'='*80}") - print(f"HYPERBOLIC DISTANCE STATISTICS") - print(f"{'='*80}") + + print(f"\n{'=' * 80}") + print("HYPERBOLIC DISTANCE STATISTICS") + print(f"{'=' * 80}") print(f"\nIntra-{rank} distances:") print(f" Count: {len(intra_distances):,}") print(f" Mean: {np.mean(intra_distances):.6f}") print(f" Median: {np.median(intra_distances):.6f}") print(f" Std: {np.std(intra_distances):.6f}") - + print(f"\nInter-{rank} distances:") print(f" Count: {len(inter_distances):,}") print(f" Mean: {np.mean(inter_distances):.6f}") print(f" Median: {np.median(inter_distances):.6f}") print(f" Std: {np.std(inter_distances):.6f}") - + # Separation separation = np.mean(inter_distances) / np.mean(intra_distances) print(f"\nSeparation Ratio: {separation:.3f}x") - + if separation > 2.0: quality = "✅ EXCELLENT" elif separation > 1.5: @@ -375,63 +387,81 @@ def analyze_hierarchical_clustering_hyperbolic(emb, idx2tax, taxonomy, rank="phy quality = "⚠️ MODERATE" else: quality = "❌ POOR" - + print(f"Quality: {quality}") - + # Plot fig, ax = plt.subplots(figsize=(12, 6)) bins = np.linspace(0, max(np.max(intra_distances), np.max(inter_distances)), 50) - ax.hist(intra_distances, bins=bins, alpha=0.6, label=f'Intra-{rank}', color='blue', density=True) - ax.hist(inter_distances, bins=bins, alpha=0.6, label=f'Inter-{rank}', color='red', density=True) - ax.axvline(np.mean(intra_distances), color='blue', linestyle='--', linewidth=2, - label=f'Intra mean: {np.mean(intra_distances):.3f}') - ax.axvline(np.mean(inter_distances), color='red', linestyle='--', linewidth=2, - label=f'Inter mean: {np.mean(inter_distances):.3f}') - ax.set_xlabel('Poincaré Distance (Hyperbolic)', fontsize=12) - ax.set_ylabel('Density', fontsize=12) - ax.set_title(f'{rank.capitalize()}-level Clustering (HYPERBOLIC)\nSeparation: {separation:.2f}x - {quality}', - fontsize=14, fontweight='bold') + ax.hist( + intra_distances, bins=bins, alpha=0.6, label=f"Intra-{rank}", color="blue", density=True + ) + ax.hist(inter_distances, bins=bins, alpha=0.6, label=f"Inter-{rank}", color="red", density=True) + ax.axvline( + np.mean(intra_distances), + color="blue", + linestyle="--", + linewidth=2, + label=f"Intra mean: {np.mean(intra_distances):.3f}", + ) + ax.axvline( + np.mean(inter_distances), + color="red", + linestyle="--", + linewidth=2, + label=f"Inter mean: {np.mean(inter_distances):.3f}", + ) + ax.set_xlabel("Poincaré Distance (Hyperbolic)", fontsize=12) + ax.set_ylabel("Density", fontsize=12) + ax.set_title( + f"{rank.capitalize()}-level Clustering (HYPERBOLIC)\nSeparation: {separation:.2f}x - {quality}", + fontsize=14, + fontweight="bold", + ) ax.legend() ax.grid(True, alpha=0.3) - + plt.tight_layout() - output_file = f'hierarchy_hyperbolic_{rank}.png' - plt.savefig(output_file, dpi=150, bbox_inches='tight') + output_file = f"hierarchy_hyperbolic_{rank}.png" + plt.savefig(output_file, dpi=150, bbox_inches="tight") print(f"\nSaved plot: {output_file}") plt.close() - + return separation, quality def main(): checkpoint = "taxonomy_model_hierarchical_small_v3_best.pth" mapping_file = "data/taxonomy_edges_small.mapping.tsv" - + # Load data emb = load_embeddings(checkpoint) idx2tax = load_mapping(mapping_file) - - valid_taxids = set(int(t) for t in idx2tax.values()) + + valid_taxids = {int(t) for t in idx2tax.values()} taxonomy = load_taxonomy_with_depth(valid_taxids) - + # 1. Check radial monotonicity pearson_r, spearman_r = analyze_depth_vs_norm(emb, idx2tax, taxonomy) - + # 2. Analyze hierarchy with hyperbolic distance results = {} for rank in ["phylum", "class", "order"]: try: - sep, qual = analyze_hierarchical_clustering_hyperbolic(emb, idx2tax, taxonomy, rank=rank) + sep, qual = analyze_hierarchical_clustering_hyperbolic( + emb, idx2tax, taxonomy, rank=rank + ) results[rank] = {"separation": sep, "quality": qual} except Exception as e: print(f"\nError analyzing {rank}: {e}") import traceback + traceback.print_exc() - + # Summary - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print("SUMMARY - HYPERBOLIC ANALYSIS") - print(f"{'='*80}\n") + print(f"{'=' * 80}\n") print(f"Depth-Norm Correlation: {pearson_r:+.3f} (Pearson)") print(f"\n{'Rank':<12} {'Separation':>12} {'Quality':>15}") print("-" * 40) diff --git a/src/taxembed/builders/__init__.py b/src/taxembed/builders/__init__.py index 3639922..4fe52a9 100644 --- a/src/taxembed/builders/__init__.py +++ b/src/taxembed/builders/__init__.py @@ -1,6 +1,5 @@ """Dataset builders for custom taxonomy slices.""" -from .taxopy_clade import build_clade_dataset, CladeBuildResult +from .taxopy_clade import CladeBuildResult, build_clade_dataset __all__ = ["build_clade_dataset", "CladeBuildResult"] - diff --git a/src/taxembed/builders/taxopy_clade.py b/src/taxembed/builders/taxopy_clade.py index f29e956..0275d95 100644 --- a/src/taxembed/builders/taxopy_clade.py +++ b/src/taxembed/builders/taxopy_clade.py @@ -6,9 +6,9 @@ import pickle import re from collections import defaultdict, deque +from collections.abc import Sequence from dataclasses import dataclass from pathlib import Path -from typing import Dict, List, Optional, Sequence, Tuple import pandas as pd import taxopy @@ -17,7 +17,7 @@ from taxembed.utils.data_validation import coverage_from_indices TaxId = int -Edge = Tuple[TaxId, TaxId] +Edge = tuple[TaxId, TaxId] @dataclass @@ -31,7 +31,7 @@ class CladeBuildResult: max_depth: int pairs_count: int output_dir: Path - files: Dict[str, Path] + files: dict[str, Path] def slugify(text: str) -> str: @@ -41,8 +41,8 @@ def slugify(text: str) -> str: return slug or "clade" -def _build_children_index(parent_map: Dict[int, int]) -> Dict[int, List[int]]: - children: Dict[int, List[int]] = defaultdict(list) +def _build_children_index(parent_map: dict[int, int]) -> dict[int, list[int]]: + children: dict[int, list[int]] = defaultdict(list) for child, parent in parent_map.items(): if child == parent: continue @@ -51,14 +51,14 @@ def _build_children_index(parent_map: Dict[int, int]) -> Dict[int, List[int]]: def _collect_clade( - children_map: Dict[int, List[int]], + children_map: dict[int, list[int]], root_taxid: int, - max_depth: Optional[int] = None, -) -> Tuple[Dict[int, int], List[Edge]]: - depth_by_taxid: Dict[int, int] = {} - edges: List[Edge] = [] + max_depth: int | None = None, +) -> tuple[dict[int, int], list[Edge]]: + depth_by_taxid: dict[int, int] = {} + edges: list[Edge] = [] - queue: deque[Tuple[int, int]] = deque([(root_taxid, 0)]) + queue: deque[tuple[int, int]] = deque([(root_taxid, 0)]) visited: set[int] = set() while queue: @@ -78,7 +78,7 @@ def _collect_clade( return depth_by_taxid, edges -def _build_mapping(depths: Dict[int, int]) -> Tuple[pd.DataFrame, Dict[int, int], Dict[int, int]]: +def _build_mapping(depths: dict[int, int]) -> tuple[pd.DataFrame, dict[int, int], dict[int, int]]: sorted_taxids = sorted(depths.keys()) mapping_df = pd.DataFrame( { @@ -86,17 +86,17 @@ def _build_mapping(depths: Dict[int, int]) -> Tuple[pd.DataFrame, Dict[int, int] "idx": list(range(len(sorted_taxids))), } ) - taxid_to_idx = dict(zip(mapping_df["taxid"], mapping_df["idx"])) + taxid_to_idx = dict(zip(mapping_df["taxid"], mapping_df["idx"], strict=False)) idx_to_taxid = {idx: taxid for taxid, idx in taxid_to_idx.items()} return mapping_df, taxid_to_idx, idx_to_taxid def _build_transitive_pairs( - depths: Dict[int, int], - parent_map: Dict[int, int], - taxid_to_idx: Dict[int, int], -) -> List[Dict[str, int]]: - pairs: List[Dict[str, int]] = [] + depths: dict[int, int], + parent_map: dict[int, int], + taxid_to_idx: dict[int, int], +) -> list[dict[str, int]]: + pairs: list[dict[str, int]] = [] iterator = depths.items() for taxid, depth in tqdm(iterator, desc="Building ancestor-descendant pairs", unit="node"): @@ -122,13 +122,13 @@ def _build_transitive_pairs( def _ensure_coverage( - pairs: List[Dict[str, int]], + pairs: list[dict[str, int]], mapping_df: pd.DataFrame, - depths: Dict[int, int], - parent_map: Dict[int, int], + depths: dict[int, int], + parent_map: dict[int, int], ) -> None: - idx_to_taxid = dict(zip(mapping_df["idx"], mapping_df["taxid"])) - taxid_to_idx = dict(zip(mapping_df["taxid"], mapping_df["idx"])) + idx_to_taxid = dict(zip(mapping_df["idx"], mapping_df["taxid"], strict=False)) + taxid_to_idx = dict(zip(mapping_df["taxid"], mapping_df["idx"], strict=False)) covered = {entry["ancestor_idx"] for entry in pairs} covered.update(entry["descendant_idx"] for entry in pairs) @@ -168,7 +168,7 @@ def _write_edges(edges: Sequence[Edge], path: Path) -> None: handle.write(f"{parent} {child}\n") -def _write_mapped_edges(edges: Sequence[Edge], mapping: Dict[int, int], path: Path) -> None: +def _write_mapped_edges(edges: Sequence[Edge], mapping: dict[int, int], path: Path) -> None: with path.open("w") as handle: for parent, child in edges: if parent in mapping and child in mapping: @@ -179,7 +179,7 @@ def _write_mapping(mapping_df: pd.DataFrame, path: Path) -> None: mapping_df.to_csv(path, sep="\t", index=False) -def _write_transitive(training_pairs: List[Dict[str, int]], prefix: Path) -> Dict[str, Path]: +def _write_transitive(training_pairs: list[dict[str, int]], prefix: Path) -> dict[str, Path]: tsv_path = prefix.with_name(f"{prefix.name}.tsv") edgelist_path = prefix.with_name(f"{prefix.name}.edgelist") pkl_path = prefix.with_name(f"{prefix.name}.pkl") @@ -211,7 +211,7 @@ def _write_manifest( edge_count: int, max_depth: int, pairs_count: int, - max_depth_requested: Optional[int], + max_depth_requested: int | None, ) -> None: manifest = { "dataset_name": dataset_name, @@ -229,10 +229,10 @@ def _write_manifest( def build_clade_dataset( root_taxid: int, *, - dataset_name: Optional[str] = None, + dataset_name: str | None = None, output_dir: Path | str = Path("data") / "taxopy", taxdump_dir: Path | str = Path("data"), - max_depth: Optional[int] = None, + max_depth: int | None = None, ) -> CladeBuildResult: """Materialize a dataset for ``root_taxid`` and its descendants.""" @@ -252,10 +252,14 @@ def build_clade_dataset( mapping_df, taxid_to_idx, idx_to_taxid = _build_mapping(depths) if not dataset_name: - root_name = taxdb.taxid2name.get(str(root_taxid)) or taxdb.taxid2name.get(root_taxid, str(root_taxid)) + root_name = taxdb.taxid2name.get(str(root_taxid)) or taxdb.taxid2name.get( + root_taxid, str(root_taxid) + ) dataset_name = slugify(root_name) else: - root_name = taxdb.taxid2name.get(str(root_taxid)) or taxdb.taxid2name.get(root_taxid, str(root_taxid)) + root_name = taxdb.taxid2name.get(str(root_taxid)) or taxdb.taxid2name.get( + root_taxid, str(root_taxid) + ) dataset_dir = output_dir / dataset_name dataset_dir.mkdir(parents=True, exist_ok=True) @@ -272,9 +276,9 @@ def build_clade_dataset( training_pairs = _build_transitive_pairs(depths, parent_map, taxid_to_idx) _ensure_coverage(training_pairs, mapping_df, depths, parent_map) - used_indices = { - entry["ancestor_idx"] for entry in training_pairs - } | {entry["descendant_idx"] for entry in training_pairs} + used_indices = {entry["ancestor_idx"] for entry in training_pairs} | { + entry["descendant_idx"] for entry in training_pairs + } coverage = coverage_from_indices(mapping_df, used_indices) if not coverage.is_perfect: missing_taxids = [idx_to_taxid[idx] for idx in sorted(coverage.missing_indices)] @@ -283,7 +287,9 @@ def build_clade_dataset( f"Missing indices: {missing_taxids[:10]}" ) - transitive_paths = _write_transitive(training_pairs, prefix.with_name(f"{prefix.name}_transitive")) + transitive_paths = _write_transitive( + training_pairs, prefix.with_name(f"{prefix.name}_transitive") + ) manifest_path = prefix.with_name(f"{prefix.name}_manifest.json") _write_manifest( @@ -316,4 +322,3 @@ def build_clade_dataset( output_dir=dataset_dir, files=files, ) - diff --git a/src/taxembed/cli/analyze.py b/src/taxembed/cli/analyze.py new file mode 100644 index 0000000..1502f3c --- /dev/null +++ b/src/taxembed/cli/analyze.py @@ -0,0 +1,13 @@ +#!/usr/bin/env python3 +"""Analyze hierarchy quality of embeddings.""" + +from taxembed.analysis import analyze_hierarchy + + +def main(): + """Analyze hierarchy quality of trained embeddings.""" + analyze_hierarchy() + + +if __name__ == "__main__": + main() diff --git a/src/taxembed/cli/check.py b/src/taxembed/cli/check.py index a2dbdf0..9f4ff0f 100644 --- a/src/taxembed/cli/check.py +++ b/src/taxembed/cli/check.py @@ -1,18 +1,12 @@ #!/usr/bin/env python3 """Check and validate trained models.""" -import sys -import os - -# Add root to path to import existing script -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))) - -from final_sanity_check import main as check_main +from taxembed.validation import run_checks def main(): """Run comprehensive sanity checks.""" - check_main() + run_checks() if __name__ == "__main__": diff --git a/src/taxembed/cli/download.py b/src/taxembed/cli/download.py index 4f1c5c8..d50148b 100644 --- a/src/taxembed/cli/download.py +++ b/src/taxembed/cli/download.py @@ -1,18 +1,12 @@ #!/usr/bin/env python3 """Download and prepare NCBI taxonomy data.""" -import sys -import os - -# Add root to path to import existing script -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))) - -from prepare_taxonomy_data import main as prepare_main +from taxembed.data import download_taxonomy def main(): """Download NCBI taxonomy data.""" - prepare_main() + download_taxonomy() if __name__ == "__main__": diff --git a/src/taxembed/cli/main.py b/src/taxembed/cli/main.py index aa2f9d5..6f65d73 100644 --- a/src/taxembed/cli/main.py +++ b/src/taxembed/cli/main.py @@ -7,17 +7,16 @@ import re import subprocess import sys -from datetime import datetime, timezone +from datetime import UTC, datetime from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any -import torch import taxopy +import torch from taxembed.builders import build_clade_dataset - -PROJECT_ROOT = Path(__file__).resolve().parents[3] # .../poincare-embeddings/src/taxembed/cli -> repo root +PROJECT_ROOT = Path(__file__).resolve().parents[3] # .../taxembed/src/taxembed/cli -> repo root DATA_DIR = PROJECT_ROOT / "data" ARTIFACTS_DIR = PROJECT_ROOT / "artifacts" / "tags" @@ -68,7 +67,9 @@ def resolve_taxid(identifier: str, taxdump_dir: Path) -> tuple[int, str]: try: taxid = int(choices[0]) except (TypeError, ValueError) as exc: # pragma: no cover - raise SystemExit(f"❌ Failed to interpret TaxID for '{identifier}': {choices[0]!r}") from exc + raise SystemExit( + f"❌ Failed to interpret TaxID for '{identifier}': {choices[0]!r}" + ) from exc name = taxdb.taxid2name.get(str(taxid)) or taxdb.taxid2name.get(taxid, str(taxid)) return taxid, name @@ -87,7 +88,7 @@ def handle_train(args: argparse.Namespace) -> None: tag_dir = ARTIFACTS_DIR / slug tag_dir.mkdir(parents=True, exist_ok=True) - dataset_record: Dict[str, Any] + dataset_record: dict[str, Any] training_data_path: Path mapping_path: Path @@ -138,11 +139,12 @@ def handle_train(args: argparse.Namespace) -> None: ) checkpoint_base = tag_dir / f"{slug}.pth" - train_script = PROJECT_ROOT / "train_small.py" + # Use the installed taxembed-train command from the package train_cmd = [ sys.executable, - str(train_script), + "-m", + "taxembed.cli.train", "--data", str(training_data_path), "--mapping", @@ -193,7 +195,7 @@ def handle_train(args: argparse.Namespace) -> None: metadata = { "tag": args.as_tag, "slug": slug, - "created_at": datetime.now(timezone.utc).isoformat(), + "created_at": datetime.now(UTC).isoformat(), "identifier": args.identifier, "dataset": dataset_record, "training": { @@ -241,19 +243,19 @@ def handle_visualize(args: argparse.Namespace) -> None: dataset_meta = metadata.get("dataset", {}) paths = metadata.get("training", {}).get("paths", {}) - checkpoint_path_str = args.checkpoint or paths.get("best_checkpoint") or paths.get("checkpoint_base", "") + checkpoint_path_str = ( + args.checkpoint or paths.get("best_checkpoint") or paths.get("checkpoint_base", "") + ) if not checkpoint_path_str: raise SystemExit(f"❌ No checkpoint path found in metadata for tag '{args.tag}'") - + checkpoint_path = Path(checkpoint_path_str) if not checkpoint_path.is_absolute(): checkpoint_path = (tag_dir / checkpoint_path_str).resolve() else: checkpoint_path = checkpoint_path.resolve() - mapping_path = Path( - args.mapping or paths.get("mapping", "") - ).resolve() + mapping_path = Path(args.mapping or paths.get("mapping", "")).resolve() if not checkpoint_path.exists(): raise SystemExit(f"❌ Checkpoint not found: {checkpoint_path}") @@ -262,10 +264,11 @@ def handle_visualize(args: argparse.Namespace) -> None: output_path = Path(args.output) if args.output else tag_dir / f"{slug}_umap.png" - viz_script = PROJECT_ROOT / "visualize_multi_groups.py" + # Use the installed taxembed visualization module viz_cmd = [ sys.executable, - str(viz_script), + "-m", + "taxembed.visualization.umap_viz", str(checkpoint_path), "--mapping", str(mapping_path), @@ -283,19 +286,21 @@ def handle_visualize(args: argparse.Namespace) -> None: viz_cmd.extend(["--root-taxid", str(root_taxid)]) # Always pass children depth (default is 0 for immediate children) viz_cmd.extend(["--children", str(args.children)]) - + # Extract title information from checkpoint and metadata - clade_name = dataset_meta.get("root_name") or dataset_meta.get("dataset_name") or args.tag.title() + clade_name = ( + dataset_meta.get("root_name") or dataset_meta.get("dataset_name") or args.tag.title() + ) epoch = None loss = None - + try: ckpt = torch.load(checkpoint_path, map_location="cpu") epoch = ckpt.get("epoch", None) loss = ckpt.get("loss", None) except Exception: pass - + if clade_name: viz_cmd.extend(["--clade-name", str(clade_name)]) if epoch is not None: @@ -313,15 +318,26 @@ def handle_visualize(args: argparse.Namespace) -> None: def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(prog="taxembed", description="Unified CLI for taxonomy embeddings") + parser = argparse.ArgumentParser( + prog="taxembed", description="Unified CLI for taxonomy embeddings" + ) subparsers = parser.add_subparsers(dest="command") train_parser = subparsers.add_parser("train", help="Build dataset and train embeddings") - train_parser.add_argument("identifier", nargs="?", help="TaxID or clade name recognized by NCBI") - train_parser.add_argument("-as", "--as-tag", required=True, help="Tag name used to reference this run") + train_parser.add_argument( + "identifier", nargs="?", help="TaxID or clade name recognized by NCBI" + ) + train_parser.add_argument( + "-as", "--as-tag", required=True, help="Tag name used to reference this run" + ) train_parser.add_argument("--file", help="Path to prebuilt transitive dataset (.pkl)") train_parser.add_argument("--mapping", help="Mapping file (required with --file)") - train_parser.add_argument("--max-depth", type=int, default=None, help="Limit descendant depth when building clades") + train_parser.add_argument( + "--max-depth", + type=int, + default=None, + help="Limit descendant depth when building clades", + ) train_parser.add_argument("--epochs", type=int, default=100) train_parser.add_argument("--dim", type=int, default=10) train_parser.add_argument("--batch-size", type=int, default=64) @@ -335,21 +351,43 @@ def build_parser() -> argparse.ArgumentParser: visualize_parser = subparsers.add_parser("visualize", help="Visualize a trained tag with UMAP") visualize_parser.add_argument("tag", help="Tag used during `taxembed train ... -as TAG`") - visualize_parser.add_argument("--sample", type=int, default=25000, help="Number of points for UMAP sampling") + visualize_parser.add_argument( + "--sample", type=int, default=25000, help="Number of points for UMAP sampling" + ) visualize_parser.add_argument("--output", help="Output image path") visualize_parser.add_argument("--checkpoint", help="Override checkpoint path") visualize_parser.add_argument("--mapping", help="Override mapping path") visualize_parser.add_argument("--names", help="Override names.dmp path") visualize_parser.add_argument("--nodes", help="Override nodes.dmp path") visualize_parser.add_argument("--root-taxid", type=int, help="Override root TaxID for coloring") - visualize_parser.add_argument("--children", type=int, default=0, - help="Depth level for coloring (0=children, 1=grandchildren, 2=great-grandchildren, etc.)") + visualize_parser.add_argument( + "--children", + type=int, + default=0, + help="Depth level for coloring (0=children, 1=grandchildren, 2=great-grandchildren, etc.)", + ) visualize_parser.set_defaults(func=handle_visualize) + analyze_parser = subparsers.add_parser( + "analyze", help="Analyze hierarchy quality of embeddings" + ) + analyze_parser.add_argument("checkpoint", help="Checkpoint file to analyze") + analyze_parser.add_argument("--mapping", help="Override mapping path") + analyze_parser.set_defaults(func=handle_analyze) + return parser -def main(argv: Optional[list[str]] = None) -> None: +def handle_analyze(args: argparse.Namespace) -> None: + """Handle analyze subcommand.""" + from taxembed.analysis import analyze_hierarchy + + # For now, just call the main function + # TODO: Update to accept checkpoint argument + analyze_hierarchy() + + +def main(argv: list[str] | None = None) -> None: parser = build_parser() args = parser.parse_args(argv) if not hasattr(args, "func"): @@ -360,4 +398,3 @@ def main(argv: Optional[list[str]] = None) -> None: if __name__ == "__main__": # pragma: no cover main() - diff --git a/src/taxembed/cli/prepare.py b/src/taxembed/cli/prepare.py index 34db8cd..2e8279f 100644 --- a/src/taxembed/cli/prepare.py +++ b/src/taxembed/cli/prepare.py @@ -1,18 +1,12 @@ #!/usr/bin/env python3 """Build transitive closure for hierarchical training.""" -import sys -import os - -# Add root to path to import existing script -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))) - -from build_transitive_closure import main as build_main +from taxembed.data.transitive import main as build_transitive_closure def main(): """Build transitive closure from taxonomy edges.""" - build_main() + build_transitive_closure() if __name__ == "__main__": diff --git a/src/taxembed/cli/train.py b/src/taxembed/cli/train.py index 65275e7..5540c5a 100644 --- a/src/taxembed/cli/train.py +++ b/src/taxembed/cli/train.py @@ -1,18 +1,123 @@ #!/usr/bin/env python3 """Train hierarchical Poincaré embeddings.""" -import sys -import os +import argparse +import pickle -# Add root to path to import existing script -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))) +import pandas as pd +import torch +import torch.optim as optim -from train_small import main as train_main +from taxembed.models import HierarchicalPoincareEmbedding +from taxembed.training import HierarchicalDataLoader, train_model def main(): """Train hierarchical model on taxonomy data.""" - train_main() + parser = argparse.ArgumentParser(description="Train hierarchical Poincaré embeddings") + parser.add_argument( + "--data", + default="data/taxonomy_edges_small_transitive.pkl", + help="Training data (default: small dataset)", + ) + parser.add_argument( + "--checkpoint", + default="taxonomy_model_small.pth", + help="Output checkpoint path", + ) + parser.add_argument( + "--mapping", + default="data/taxonomy_edges_small.mapping.tsv", + help="Mapping file path aligned with training data", + ) + parser.add_argument("--dim", type=int, default=10, help="Embedding dimension") + parser.add_argument("--epochs", type=int, default=100, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64, help="Batch size") + parser.add_argument("--n-negatives", type=int, default=50, help="Number of negative samples") + parser.add_argument("--lr", type=float, default=0.005, help="Learning rate") + parser.add_argument("--margin", type=float, default=0.2, help="Ranking loss margin") + parser.add_argument("--lambda-reg", type=float, default=0.1, help="Regularization strength") + parser.add_argument( + "--early-stopping", + type=int, + default=5, + help="Early stopping patience (0 to disable)", + ) + parser.add_argument("--gpu", type=int, default=-1, help="GPU device index (-1 for CPU)") + + args = parser.parse_args() + + # Device + if args.gpu >= 0 and torch.cuda.is_available(): + device = torch.device(f"cuda:{args.gpu}") + else: + device = torch.device("cpu") + + print("Loading training data...") + with open(args.data, "rb") as f: + training_data = pickle.load(f) + + print(f" ✓ Loaded {len(training_data):,} training pairs") + + # Load mapping to get node count + print("Loading mapping...") + mapping_df = pd.read_csv(args.mapping, sep="\t", header=None, names=["idx", "taxid"]) + n_nodes = len(mapping_df) + print(f" ✓ {n_nodes:,} unique nodes") + + # Build depth map + print("Building depth map...") + idx_to_depth = {} + max_depth = 0 + for item in training_data: + ancestor_idx = item["ancestor_idx"] + descendant_idx = item["descendant_idx"] + ancestor_depth = item["ancestor_depth"] + descendant_depth = item["descendant_depth"] + + idx_to_depth[ancestor_idx] = ancestor_depth + idx_to_depth[descendant_idx] = descendant_depth + max_depth = max(max_depth, ancestor_depth, descendant_depth) + + print(f" ✓ Depth range: [0, {max_depth}]") + + # Create model + print(f"\nInitializing {args.dim}D Poincaré embeddings...") + model = HierarchicalPoincareEmbedding( + n_nodes=n_nodes, dim=args.dim, max_depth=max_depth, init_depth_data=idx_to_depth + ) + + # Create data loader + print("Creating data loader...") + dataloader = HierarchicalDataLoader( + training_data=training_data, + n_nodes=n_nodes, + batch_size=args.batch_size, + n_negatives=args.n_negatives, + depth_stratify=True, + ) + + # Optimizer + optimizer = optim.Adam(model.parameters(), lr=args.lr) + + # Train + train_model( + model=model, + dataloader=dataloader, + optimizer=optimizer, + n_epochs=args.epochs, + idx_to_depth=idx_to_depth, + max_depth=max_depth, + device=device, + margin=args.margin, + lambda_reg=args.lambda_reg, + early_stopping_patience=args.early_stopping, + checkpoint_base=args.checkpoint, + ) + + print( + f"\n✓ Training complete! Best model saved to {args.checkpoint.replace('.pth', '_best.pth')}" + ) if __name__ == "__main__": diff --git a/src/taxembed/cli/visualize.py b/src/taxembed/cli/visualize.py index ae5b4eb..c653668 100644 --- a/src/taxembed/cli/visualize.py +++ b/src/taxembed/cli/visualize.py @@ -1,19 +1,12 @@ #!/usr/bin/env python3 """Visualize taxonomy embeddings with UMAP.""" -import sys -import os -import argparse - -# Add root to path to import existing script -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))) - -from visualize_multi_groups import main as viz_main +from taxembed.visualization import visualize_embeddings def main(): """Create UMAP visualization of embeddings.""" - viz_main() + visualize_embeddings() if __name__ == "__main__": diff --git a/src/taxembed/data/__init__.py b/src/taxembed/data/__init__.py new file mode 100644 index 0000000..7cd9d8f --- /dev/null +++ b/src/taxembed/data/__init__.py @@ -0,0 +1,17 @@ +"""Data downloading and processing utilities.""" + +from .download import ( + ensure_taxdump, + parse_nodes_dmp, + parse_names_dmp, + download_taxonomy, +) +from .transitive import build_transitive_closure + +__all__ = [ + "ensure_taxdump", + "parse_nodes_dmp", + "parse_names_dmp", + "download_taxonomy", + "build_transitive_closure", +] diff --git a/prepare_taxonomy_data.py b/src/taxembed/data/download.py similarity index 79% rename from prepare_taxonomy_data.py rename to src/taxembed/data/download.py index bedcef4..b604f2d 100644 --- a/prepare_taxonomy_data.py +++ b/src/taxembed/data/download.py @@ -13,66 +13,65 @@ import pandas as pd from tqdm import tqdm + def parse_nodes_dmp(nodes_file): """ Parse nodes.dmp file to extract tax_id and parent_tax_id. - + Format: tax_id | parent_tax_id | rank | ... Delimiter: \t|\t """ print(f"Parsing {nodes_file}...") - + edges = [] - with open(nodes_file, 'r') as f: + with open(nodes_file, "r") as f: for line in tqdm(f, desc="Reading nodes"): # Split by \t|\t delimiter - parts = line.strip().split('\t|\t') + parts = line.strip().split("\t|\t") if len(parts) >= 2: try: tax_id = int(parts[0].strip()) parent_tax_id = int(parts[1].strip()) - + # Skip self-loops (root node) if tax_id != parent_tax_id: - edges.append({ - 'id1': tax_id, - 'id2': parent_tax_id, - 'weight': 1.0 - }) + edges.append({"id1": tax_id, "id2": parent_tax_id, "weight": 1.0}) except (ValueError, IndexError): continue - + return edges + def parse_names_dmp(names_file): """ Parse names.dmp file to extract scientific names. - + Format: tax_id | name | unique_name | name_class Delimiter: \t|\t """ print(f"Parsing {names_file}...") - + names_map = {} - with open(names_file, 'r') as f: + with open(names_file, "r") as f: for line in tqdm(f, desc="Reading names"): - parts = line.strip().split('\t|\t') + parts = line.strip().split("\t|\t") if len(parts) >= 4: try: tax_id = int(parts[0].strip()) name = parts[1].strip() name_class = parts[3].strip() - + # Prefer scientific names, but keep first occurrence if tax_id not in names_map: names_map[tax_id] = name - elif name_class == 'scientific name': + elif name_class == "scientific name": names_map[tax_id] = name except (ValueError, IndexError): continue - + return names_map + TAXDUMP_URL = "https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/new_taxdump/new_taxdump.tar.gz" @@ -90,7 +89,10 @@ def ensure_taxdump(data_dir: Path) -> tuple[Path, Path, Path | None]: archive_path = data_dir / "new_taxdump.tar.gz" print(f"Downloading NCBI taxonomy → {archive_path}") - with urllib.request.urlopen(TAXDUMP_URL) as response, archive_path.open("wb") as out_f: + with ( + urllib.request.urlopen(TAXDUMP_URL) as response, + archive_path.open("wb") as out_f, + ): shutil.copyfileobj(response, out_f) print("Extracting taxdump archive...") @@ -103,52 +105,71 @@ def ensure_taxdump(data_dir: Path) -> tuple[Path, Path, Path | None]: return nodes_file, names_file, merged_file -def main(): - data_dir = Path(__file__).parent / "data" +def download_taxonomy(data_dir: Path | None = None): + """Download and parse NCBI taxonomy data. + + Args: + data_dir: Directory for data files (defaults to ./data) + + Returns: + Path to the output edgelist file + """ + if data_dir is None: + data_dir = Path.cwd() / "data" + + data_dir = Path(data_dir) nodes_file, names_file, _ = ensure_taxdump(data_dir) output_file = data_dir / "taxonomy_edges.csv" - + # Parse taxonomy data edges = parse_nodes_dmp(nodes_file) names_map = parse_names_dmp(names_file) - + print(f"\nExtracted {len(edges)} edges from taxonomy") print(f"Extracted {len(names_map)} taxonomy names") - + # Create DataFrame df = pd.DataFrame(edges) - + print(f"\nDataFrame shape: {df.shape}") print(f"Sample edges:\n{df.head(10)}") - + # Save to CSV df.to_csv(output_file, index=False) print(f"\n✓ Saved edge list to {output_file}") - + # Also create edgelist format (no header, whitespace-separated) - edgelist_file = output_file.with_suffix('.edgelist') - with open(edgelist_file, 'w') as f: + edgelist_file = output_file.with_suffix(".edgelist") + with open(edgelist_file, "w") as f: for _, row in df.iterrows(): f.write(f"{int(row['id1'])} {int(row['id2'])}\n") print(f"✓ Saved edgelist to {edgelist_file}") - + # Print statistics print(f"\nStatistics:") print(f" Total edges: {len(df)}") print(f" Unique parent nodes: {df['id2'].nunique()}") print(f" Unique child nodes: {df['id1'].nunique()}") - unique_nodes = set(df['id1'].unique()) | set(df['id2'].unique()) + unique_nodes = set(df["id1"].unique()) | set(df["id2"].unique()) print(f" Total unique nodes: {len(unique_nodes)}") - + # Show some example edges with names print(f"\nExample edges with names:") sample_edges = df.head(10) for _, row in sample_edges.iterrows(): - child_id = int(row['id1']) - parent_id = int(row['id2']) + child_id = int(row["id1"]) + parent_id = int(row["id2"]) child_name = names_map.get(child_id, f"Unknown_{child_id}") parent_name = names_map.get(parent_id, f"Unknown_{parent_id}") print(f" {child_name} (id={child_id}) → {parent_name} (id={parent_id})") -if __name__ == '__main__': + return edgelist_file + + +def main(): + """CLI entry point.""" + download_taxonomy() + + +if __name__ == "__main__": main() diff --git a/remap_edges.py b/src/taxembed/data/mapping.py similarity index 100% rename from remap_edges.py rename to src/taxembed/data/mapping.py diff --git a/build_transitive_closure.py b/src/taxembed/data/transitive.py similarity index 100% rename from build_transitive_closure.py rename to src/taxembed/data/transitive.py diff --git a/src/taxembed/models/__init__.py b/src/taxembed/models/__init__.py index 8fa8564..ff5717b 100644 --- a/src/taxembed/models/__init__.py +++ b/src/taxembed/models/__init__.py @@ -1 +1,9 @@ -"""Embedding models.""" +"""Embedding models and utilities.""" + +from .metrics import MetricsTracker +from .poincare import HierarchicalPoincareEmbedding + +__all__ = [ + "HierarchicalPoincareEmbedding", + "MetricsTracker", +] diff --git a/src/taxembed/models/metrics.py b/src/taxembed/models/metrics.py new file mode 100644 index 0000000..7d930cf --- /dev/null +++ b/src/taxembed/models/metrics.py @@ -0,0 +1,121 @@ +"""Training metrics tracker with visualization.""" + + +class MetricsTracker: + """Track and display training metrics with visual improvements.""" + + def __init__(self): + self.history = [] + self.best_loss = float("inf") + self.best_epoch = 0 + + def update(self, epoch: int, metrics: dict): + """Update metrics for current epoch. + + Args: + epoch: Current epoch number + metrics: Dictionary of metric values + """ + self.history.append(metrics) + + if metrics["loss"] < self.best_loss: + self.best_loss = metrics["loss"] + self.best_epoch = epoch + + def get_previous(self, metric_name: str): + """Get previous epoch's metric value. + + Args: + metric_name: Name of the metric to retrieve + + Returns: + Previous metric value or None if not available + """ + if len(self.history) < 2: + return None + return self.history[-2].get(metric_name) + + def print_header(self): + """Print column headers for metrics table.""" + print("\n" + "=" * 100) + print( + f"{'Epoch':>6} | {'Loss':>10} | {'ΔLoss':>10} | {'Improve':>8} | " + f"{'Reg':>8} | {'MaxNorm':>8} | {'Outside':>7} | {'Status':>10}" + ) + print("=" * 100) + + def print_epoch_summary(self, epoch: int, metrics: dict, total_epochs: int): + """Print compact summary of epoch with improvement indicators. + + Args: + epoch: Current epoch number + metrics: Dictionary of current metrics + total_epochs: Total number of epochs planned + """ + prev_loss = self.get_previous("loss") + + # Calculate improvement + if prev_loss is not None: + delta = metrics["loss"] - prev_loss + pct_change = (delta / prev_loss) * 100 if prev_loss != 0 else 0 + + if delta < 0: + status = "✓ BETTER" + delta_str = f"{delta:+.4f}" + pct_str = f"{pct_change:+.2f}%" + improve_color = "\033[92m" # Green + else: + status = "✗ WORSE" + delta_str = f"{delta:+.4f}" + pct_str = f"{pct_change:+.2f}%" + improve_color = "\033[91m" # Red + + reset_color = "\033[0m" + else: + delta_str = "---" + pct_str = "---" + status = "FIRST" + improve_color = "" + reset_color = "" + + # Format output + outside_pct = ( + (metrics["outside_count"] / metrics["total_nodes"]) * 100 + if metrics["total_nodes"] > 0 + else 0 + ) + + print( + f"{epoch:6d} | " + f"{metrics['loss']:10.6f} | " + f"{improve_color}{delta_str:>10}{reset_color} | " + f"{improve_color}{pct_str:>8}{reset_color} | " + f"{metrics['reg_loss']:8.6f} | " + f"{metrics['max_norm']:8.4f} | " + f"{outside_pct:6.2f}% | " + f"{improve_color}{status:>10}{reset_color}" + ) + + # Additional info every 5 epochs + if epoch % 5 == 0 or epoch == 1: + print( + f" └─ Best: {self.best_loss:.6f} @ epoch {self.best_epoch} | " + f"Norms: [{metrics['min_norm']:.4f}, {metrics['mean_norm']:.4f}, {metrics['max_norm']:.4f}]" + ) + + def print_final_summary(self): + """Print final training summary with overall statistics.""" + print("\n" + "=" * 100) + print("TRAINING SUMMARY") + print("=" * 100) + print(f"Total epochs: {len(self.history)}") + print(f"Best loss: {self.best_loss:.6f} (epoch {self.best_epoch})") + + if len(self.history) >= 2: + first_loss = self.history[0]["loss"] + last_loss = self.history[-1]["loss"] + total_improvement = first_loss - last_loss + pct_improvement = (total_improvement / first_loss) * 100 + print(f"Total improvement: {total_improvement:+.6f} ({pct_improvement:+.2f}%)") + + print("=" * 100 + "\n") diff --git a/src/taxembed/models/poincare.py b/src/taxembed/models/poincare.py new file mode 100644 index 0000000..75ce895 --- /dev/null +++ b/src/taxembed/models/poincare.py @@ -0,0 +1,149 @@ +"""Poincaré embeddings model with hierarchical features.""" + +import torch +import torch.nn as nn + + +class HierarchicalPoincareEmbedding(nn.Module): + """Poincaré embeddings with hierarchical structure. + + Key features: + - Depth-aware initialization (deeper nodes near boundary) + - Hyperbolic distance computation + - Ball constraint projection + + Args: + n_nodes: Number of nodes to embed + dim: Embedding dimensionality + max_depth: Maximum depth in hierarchy (for initialization) + init_depth_data: Optional dict mapping node idx -> depth for initialization + """ + + def __init__( + self, + n_nodes: int, + dim: int = 10, + max_depth: int = 38, + init_depth_data: dict | None = None, + ): + super().__init__() + self.n_nodes = n_nodes + self.dim = dim + self.max_depth = max_depth + + # Embeddings (initialize later with depth info) + self.embeddings = nn.Embedding(n_nodes, dim) + + # Initialize based on depth if available + if init_depth_data is not None: + self._initialize_by_depth(init_depth_data) + else: + # Default: uniform small initialization + nn.init.uniform_(self.embeddings.weight, -0.001, 0.001) + + def _initialize_by_depth(self, depth_data: dict): + """Initialize embeddings based on taxonomic depth. + + Deeper nodes → larger radius (closer to boundary). + This encodes hierarchy from the start! + + Args: + depth_data: Dict mapping idx -> depth + """ + print("Initializing embeddings by depth...") + + # Map: idx → depth + idx_to_depth = depth_data + + with torch.no_grad(): + for idx in range(self.n_nodes): + depth = idx_to_depth.get(idx, 0) + + # Radius increases with depth + # Root (depth 0): r ≈ 0.1 + # Max depth: r ≈ 0.95 + target_radius = 0.1 + (depth / self.max_depth) * 0.85 + + # Random direction on sphere + vec = torch.randn(self.dim) + vec = vec / vec.norm() + + # Scale to target radius + self.embeddings.weight[idx] = vec * target_radius + + norms = self.embeddings.weight.norm(dim=1) + print(f" ✓ Initialized: norm range [{norms.min():.3f}, {norms.max():.3f}]") + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + """Get embeddings for indices. + + Args: + indices: Tensor of node indices + + Returns: + Embeddings tensor of shape (batch_size, dim) + """ + return self.embeddings(indices) + + def poincare_distance( + self, u: torch.Tensor, v: torch.Tensor, eps: float = 1e-5 + ) -> torch.Tensor: + """Compute Poincaré distance between embeddings. + + Formula: d(u,v) = arcosh(1 + 2||u-v||²/((1-||u||²)(1-||v||²))) + + Args: + u: First embedding tensor + v: Second embedding tensor + eps: Small value for numerical stability + + Returns: + Poincaré distances + """ + # Compute squared norms + u_norm_sq = (u**2).sum(dim=-1) + v_norm_sq = (v**2).sum(dim=-1) + + # Clamp to stay inside ball + u_norm_sq = torch.clamp(u_norm_sq, 0, 1 - eps) + v_norm_sq = torch.clamp(v_norm_sq, 0, 1 - eps) + + # Squared Euclidean distance + diff_norm_sq = ((u - v) ** 2).sum(dim=-1) + + # Poincaré distance + numerator = 2 * diff_norm_sq + denominator = (1 - u_norm_sq) * (1 - v_norm_sq) + + dist = torch.acosh(1 + numerator / (denominator + eps) + eps) + + return dist + + def project_to_ball(self, indices: torch.Tensor | None = None, max_norm: float = 0.999): + """Project embeddings back into Poincaré ball with HARD constraint. + + Args: + indices: If provided, only project these indices (more efficient). + If None, project all embeddings. + max_norm: Maximum allowed norm (default 0.999, essentially at boundary) + """ + with torch.no_grad(): + if indices is not None: + # Only project updated embeddings + embs = self.embeddings.weight[indices] + norms = embs.norm(dim=1, keepdim=True) + # Hard projection: if norm >= max_norm, scale it down + # Use where to only scale embeddings that need it + needs_projection = norms >= max_norm + scale = torch.where( + needs_projection, max_norm / (norms + 1e-8), torch.ones_like(norms) + ) + self.embeddings.weight[indices] = embs * scale + else: + # Project all embeddings + norms = self.embeddings.weight.norm(dim=1, keepdim=True) + needs_projection = norms >= max_norm + scale = torch.where( + needs_projection, max_norm / (norms + 1e-8), torch.ones_like(norms) + ) + self.embeddings.weight.mul_(scale) diff --git a/src/taxembed/training/__init__.py b/src/taxembed/training/__init__.py new file mode 100644 index 0000000..ee6b552 --- /dev/null +++ b/src/taxembed/training/__init__.py @@ -0,0 +1,12 @@ +"""Training utilities and data loaders.""" + +from .data_loader import HierarchicalDataLoader +from .loss import radial_regularizer, ranking_loss_with_margin +from .trainer import train_model + +__all__ = [ + "HierarchicalDataLoader", + "ranking_loss_with_margin", + "radial_regularizer", + "train_model", +] diff --git a/src/taxembed/training/data_loader.py b/src/taxembed/training/data_loader.py new file mode 100644 index 0000000..3b95423 --- /dev/null +++ b/src/taxembed/training/data_loader.py @@ -0,0 +1,130 @@ +"""Hierarchical data loader with depth-aware sampling.""" + +from collections import defaultdict +from collections.abc import Iterator + +import numpy as np +import torch + + +class HierarchicalDataLoader: + """Data loader with depth-aware sampling and hard negatives. + + Args: + training_data: List of dicts with metadata (ancestor_idx, descendant_idx, etc.) + n_nodes: Total number of nodes + batch_size: Batch size for training + n_negatives: Number of negative samples per positive pair + depth_stratify: Whether to use depth-aware stratification + """ + + def __init__( + self, + training_data: list[dict], + n_nodes: int, + batch_size: int = 32, + n_negatives: int = 50, + depth_stratify: bool = True, + ): + self.training_data = training_data + self.n_nodes = n_nodes + self.batch_size = batch_size + self.n_negatives = n_negatives + self.depth_stratify = depth_stratify + + # Build index by depth for stratified sampling + if depth_stratify: + self.depth_buckets = defaultdict(list) + for i, item in enumerate(training_data): + depth_diff = item["depth_diff"] + self.depth_buckets[depth_diff].append(i) + print(f" ✓ Created {len(self.depth_buckets)} depth buckets for sampling") + + # Build node → siblings map for hard negatives + self._build_sibling_map() + + def _build_sibling_map(self): + """Build map: node → nodes at same depth (for hard negatives).""" + print(" Building sibling map for hard negatives...") + + depth_to_nodes = defaultdict(set) + for item in self.training_data: + depth_to_nodes[item["descendant_depth"]].add(item["descendant_idx"]) + + self.sibling_map = {} + for _depth, nodes in depth_to_nodes.items(): + nodes_list = list(nodes) + for node in nodes_list: + # Siblings = other nodes at same depth + self.sibling_map[node] = [n for n in nodes_list if n != node] + + print(" ✓ Built sibling map for hard negatives") + + def __len__(self) -> int: + return len(self.training_data) // self.batch_size + + def __iter__( + self, + ) -> Iterator[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]: + """Iterate over batches with depth-aware sampling. + + Yields: + Tuple of (ancestors, descendants, negatives, depths) tensors + """ + indices = list(range(len(self.training_data))) + + if self.depth_stratify: + # Stratified sampling: mix shallow and deep pairs + np.random.shuffle(indices) + else: + # Random sampling + np.random.shuffle(indices) + + for i in range(0, len(indices), self.batch_size): + batch_indices = indices[i : i + self.batch_size] + batch = [self.training_data[idx] for idx in batch_indices] + + # Extract data (vectorized for speed) + batch_size = len(batch) + ancestors_np = np.zeros(batch_size, dtype=np.int64) + descendants_np = np.zeros(batch_size, dtype=np.int64) + depths_np = np.zeros(batch_size, dtype=np.float32) + + for j, item in enumerate(batch): + ancestors_np[j] = item["ancestor_idx"] + descendants_np[j] = item["descendant_idx"] + depths_np[j] = item["depth_diff"] + + # Convert to tensors once + ancestors = torch.from_numpy(ancestors_np) + descendants = torch.from_numpy(descendants_np) + depths = torch.from_numpy(depths_np) + + # Sample hard negatives (cousins at same depth) + negatives_np = np.zeros((batch_size, self.n_negatives), dtype=np.int64) + for j, item in enumerate(batch): + desc_idx = item["descendant_idx"] + + # Get siblings (nodes at same depth) + siblings = self.sibling_map.get(desc_idx, []) + + if len(siblings) >= self.n_negatives: + # Sample from siblings (hard negatives) + negatives_np[j] = np.random.choice(siblings, self.n_negatives, replace=False) + else: + # Mix siblings + random negatives + n_sibling = len(siblings) + n_random = self.n_negatives - n_sibling + if n_sibling > 0: + negatives_np[j, :n_sibling] = siblings + negatives_np[j, n_sibling:] = np.random.choice( + self.n_nodes, n_random, replace=False + ) + else: + negatives_np[j] = np.random.choice( + self.n_nodes, self.n_negatives, replace=False + ) + + negatives = torch.from_numpy(negatives_np) + + yield ancestors, descendants, negatives, depths diff --git a/src/taxembed/training/loss.py b/src/taxembed/training/loss.py new file mode 100644 index 0000000..cf858c5 --- /dev/null +++ b/src/taxembed/training/loss.py @@ -0,0 +1,92 @@ +"""Loss functions for hierarchical Poincaré embeddings.""" + +import torch + + +def ranking_loss_with_margin( + model, + ancestors: torch.Tensor, + descendants: torch.Tensor, + negatives: torch.Tensor, + depths: torch.Tensor, + margin: float = 0.1, + depth_weight: bool = True, +) -> torch.Tensor: + """Ranking loss with margin and optional depth weighting. + + Loss encourages: + - d(ancestor, descendant) < d(ancestor, negative) + margin + - Deeper pairs get higher weight (they're more informative) + + Args: + model: Poincaré embedding model + ancestors: Tensor of ancestor indices (batch,) + descendants: Tensor of descendant indices (batch,) + negatives: Tensor of negative sample indices (batch, n_neg) + depths: Tensor of depth differences (batch,) + margin: Margin for ranking loss + depth_weight: Whether to weight loss by depth + + Returns: + Scalar loss value + """ + # Get embeddings + anc_emb = model(ancestors) # (batch, dim) + desc_emb = model(descendants) # (batch, dim) + neg_emb = model(negatives) # (batch, n_neg, dim) + + # Positive distances (ancestor → descendant) + pos_dist = model.poincare_distance(anc_emb, desc_emb) # (batch,) + + # Negative distances (ancestor → each negative) + # Expand anc_emb to match negatives shape + anc_emb_expanded = anc_emb.unsqueeze(1).expand_as(neg_emb) # (batch, n_neg, dim) + neg_dist = model.poincare_distance(anc_emb_expanded, neg_emb) # (batch, n_neg) + + # Margin ranking loss: max(0, pos_dist - neg_dist + margin) + losses = torch.relu(pos_dist.unsqueeze(1) - neg_dist + margin) # (batch, n_neg) + loss = losses.mean(dim=1) # Average over negatives: (batch,) + + # Depth weighting: deeper pairs are more important + if depth_weight: + # Weight = sqrt(depth) to emphasize deep pairs without over-weighting + weights = torch.sqrt(depths + 1) # +1 to avoid zero weight + weights = weights / weights.mean() # Normalize + loss = loss * weights + + return loss.mean() + + +def radial_regularizer( + model, + idx_to_depth_tensor: torch.Tensor, + target_radii_tensor: torch.Tensor, + lambda_reg: float = 0.01, +) -> torch.Tensor: + """Vectorized radial regularizer to keep nodes at expected radius based on depth. + + Encourages: ||embedding|| ≈ f(depth) + where f(depth) = 0.1 + (depth/max_depth) * 0.85 + + Args: + model: Poincaré embedding model + idx_to_depth_tensor: Tensor of indices to regularize + target_radii_tensor: Tensor of target radii for each index + lambda_reg: Regularization strength + + Returns: + Scalar regularization loss + """ + if len(idx_to_depth_tensor) == 0: + return torch.tensor(0.0, device=model.embeddings.weight.device) + + # Get embeddings for these indices + embs = model.embeddings.weight[idx_to_depth_tensor] # (n, dim) + + # Compute actual radii + actual_radii = embs.norm(dim=1) # (n,) + + # L2 penalty + reg_loss = ((actual_radii - target_radii_tensor) ** 2).mean() + + return lambda_reg * reg_loss diff --git a/src/taxembed/training/trainer.py b/src/taxembed/training/trainer.py new file mode 100644 index 0000000..841d8c8 --- /dev/null +++ b/src/taxembed/training/trainer.py @@ -0,0 +1,228 @@ +"""Main training loop for Poincaré embeddings.""" + +import os +from collections import deque + +import torch +from tqdm import tqdm + +from ..models import MetricsTracker +from .loss import radial_regularizer, ranking_loss_with_margin + + +def train_model( + model, + dataloader, + optimizer, + n_epochs: int, + idx_to_depth: dict[int, int], + max_depth: int, + device: torch.device, + margin: float = 0.2, + lambda_reg: float = 0.1, + early_stopping_patience: int = 3, + checkpoint_base: str | None = None, +): + """Train Poincaré embeddings with hierarchical features. + + Args: + model: HierarchicalPoincareEmbedding model + dataloader: HierarchicalDataLoader instance + optimizer: PyTorch optimizer + n_epochs: Number of training epochs + idx_to_depth: Dictionary mapping node index to depth + max_depth: Maximum depth in hierarchy + device: PyTorch device (cpu or cuda) + margin: Margin for ranking loss + lambda_reg: Regularization strength + early_stopping_patience: Patience for early stopping (0 to disable) + checkpoint_base: Base path for saving checkpoints + + Returns: + Trained model + """ + model.to(device) + model.train() + + print("\n" + "🚀 " * 20) + print("HIERARCHICAL TRAINING - POINCARÉ EMBEDDINGS") + print("🚀 " * 20) + print("\nConfiguration:") + print(f" Margin: {margin}") + print(f" Regularization: λ={lambda_reg}") + print( + f" Early stopping: {'disabled' if early_stopping_patience == 0 else f'{early_stopping_patience} epochs'}" + ) + print(f" Device: {device}") + print(f" Max epochs: {n_epochs}") + + # Precompute regularizer tensors + print("\nPrecomputing regularization targets...") + reg_indices = [] + reg_target_radii = [] + for idx, depth in idx_to_depth.items(): + if idx < model.n_nodes: + reg_indices.append(idx) + target_radius = 0.1 + (depth / max_depth) * 0.85 + reg_target_radii.append(target_radius) + + reg_indices_tensor = torch.LongTensor(reg_indices).to(device) + reg_target_radii_tensor = torch.FloatTensor(reg_target_radii).to(device) + print(f" ✓ Regularizing {len(reg_indices):,} nodes") + + # Metrics tracker + tracker = MetricsTracker() + tracker.print_header() + + # Early stopping + epochs_without_improvement = 0 + checkpoint_queue: deque = deque(maxlen=5) + + for epoch in range(1, n_epochs + 1): + epoch_loss = 0.0 + epoch_reg_loss = 0.0 + n_batches = 0 + + # Progress bar for batches + pbar = tqdm( + dataloader, + desc=f"Epoch {epoch:3d}/{n_epochs}", + bar_format="{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]", + ncols=100, + leave=False, + dynamic_ncols=True, + mininterval=0.5, + ) + + for ancestors, descendants, negatives, depths in pbar: + ancestors = ancestors.to(device) + descendants = descendants.to(device) + negatives = negatives.to(device) + depths = depths.to(device) + + optimizer.zero_grad() + + # Ranking loss + loss = ranking_loss_with_margin( + model, + ancestors, + descendants, + negatives, + depths, + margin=margin, + depth_weight=True, + ) + + # Radial regularizer + reg_loss = radial_regularizer( + model, reg_indices_tensor, reg_target_radii_tensor, lambda_reg + ) + + # Total loss + total_loss = loss + reg_loss + total_loss.backward() + + # Gradient clipping + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + + optimizer.step() + + # Selective projection + updated_indices = torch.cat([ancestors, descendants, negatives.flatten()]) + updated_indices = torch.unique(updated_indices) + model.project_to_ball(updated_indices) + + # Periodic full projection + if n_batches % 500 == 0: + model.project_to_ball(indices=None) + + epoch_loss += loss.item() + epoch_reg_loss += reg_loss.item() + n_batches += 1 + + # Update progress bar with current batch metrics + pbar.set_postfix({"loss": f"{loss.item():.4f}", "reg": f"{reg_loss.item():.4f}"}) + + # Final projection at epoch end + model.project_to_ball(indices=None) + + # Compute metrics + avg_loss = epoch_loss / n_batches + avg_reg = epoch_reg_loss / n_batches + norms = model.embeddings.weight.norm(dim=1) + outside_count = (norms >= 1.0).sum().item() + + metrics = { + "loss": avg_loss, + "reg_loss": avg_reg, + "min_norm": norms.min().item(), + "mean_norm": norms.mean().item(), + "max_norm": norms.max().item(), + "outside_count": outside_count, + "total_nodes": model.n_nodes, + } + + # Early stopping check - MUST happen BEFORE updating tracker + prev_best_loss = tracker.best_loss + + # Update tracker and display + tracker.update(epoch, metrics) + tracker.print_epoch_summary(epoch, metrics, n_epochs) + + # Save checkpoint + if checkpoint_base: + checkpoint_path = checkpoint_base.replace(".pth", f"_epoch{epoch}.pth") + torch.save( + { + "state_dict": {"lt.weight": model.embeddings.weight}, + "embeddings": model.embeddings.weight, + "epoch": epoch, + "loss": avg_loss, + "reg_loss": avg_reg, + "best_loss": tracker.best_loss, + "epochs_without_improvement": epochs_without_improvement, + }, + checkpoint_path, + ) + + # Manage checkpoint queue + if checkpoint_queue.maxlen and len(checkpoint_queue) >= checkpoint_queue.maxlen: + old_checkpoint = checkpoint_queue[0] + if os.path.exists(old_checkpoint): + os.remove(old_checkpoint) + + checkpoint_queue.append(checkpoint_path) + + # Check if loss improved (compare against PREVIOUS best, not updated best) + if avg_loss < prev_best_loss: + epochs_without_improvement = 0 + + # Save best model + if checkpoint_base: + best_checkpoint = checkpoint_base.replace(".pth", "_best.pth") + torch.save( + { + "state_dict": {"lt.weight": model.embeddings.weight}, + "embeddings": model.embeddings.weight, + "epoch": epoch, + "loss": avg_loss, + "reg_loss": avg_reg, + }, + best_checkpoint, + ) + else: + epochs_without_improvement += 1 + + # Only check early stopping if patience > 0 (0 means disabled) + if ( + early_stopping_patience > 0 + and epochs_without_improvement >= early_stopping_patience + ): + print(f"\n🛑 Early stopping triggered after {epoch} epochs") + print(f" Best loss: {tracker.best_loss:.6f} (epoch {tracker.best_epoch})") + break + + # Final summary + tracker.print_final_summary() + + return model diff --git a/src/taxembed/utils/data_validation.py b/src/taxembed/utils/data_validation.py index b837366..2836d9e 100644 --- a/src/taxembed/utils/data_validation.py +++ b/src/taxembed/utils/data_validation.py @@ -2,9 +2,9 @@ from __future__ import annotations +from collections.abc import Iterable from dataclasses import dataclass from pathlib import Path -from typing import Iterable, Set import pandas as pd @@ -15,7 +15,7 @@ class CoverageReport: total_nodes: int used_nodes: int - missing_indices: Set[int] + missing_indices: set[int] @property def missing_count(self) -> int: @@ -49,7 +49,7 @@ def load_mapping(mapping_path: Path) -> pd.DataFrame: return df -def mapping_indices(df: pd.DataFrame) -> Set[int]: +def mapping_indices(df: pd.DataFrame) -> set[int]: """Return the set of sequential indices present in a mapping dataframe.""" return set(df["idx"].astype(int).tolist()) @@ -74,4 +74,3 @@ def coverage_from_indices(df: pd.DataFrame, used_indices: Iterable[int]) -> Cove "load_mapping", "mapping_indices", ] - diff --git a/src/taxembed/validation/__init__.py b/src/taxembed/validation/__init__.py new file mode 100644 index 0000000..40f830e --- /dev/null +++ b/src/taxembed/validation/__init__.py @@ -0,0 +1,7 @@ +"""Validation and sanity checks for data and models.""" + +from .checks import main as run_checks + +__all__ = [ + "run_checks", +] diff --git a/final_sanity_check.py b/src/taxembed/validation/checks.py similarity index 82% rename from final_sanity_check.py rename to src/taxembed/validation/checks.py index 948ea5d..550da48 100644 --- a/final_sanity_check.py +++ b/src/taxembed/validation/checks.py @@ -5,10 +5,12 @@ """ import os -import sys -import torch import pickle +import sys + import pandas as pd +import torch + def check_file_exists(path, description): """Check if file exists.""" @@ -20,73 +22,76 @@ def check_file_exists(path, description): print(f" ❌ {description} MISSING: {path}") return False + def check_model(path, expected_shape): """Check model checkpoint integrity.""" try: - ckpt = torch.load(path, map_location='cpu') - embeddings = ckpt['embeddings'] - + ckpt = torch.load(path, map_location="cpu") + embeddings = ckpt["embeddings"] + # Check shape if embeddings.shape != expected_shape: print(f" ❌ Shape mismatch: {embeddings.shape} != {expected_shape}") return False - + # Check ball constraint norms = embeddings.norm(dim=1).detach().numpy() outside = (norms >= 1.0).sum() max_norm = norms.max() - + if outside > 0: print(f" ❌ {outside} embeddings outside ball") return False - + if max_norm > 1.0: print(f" ❌ Max norm {max_norm:.4f} > 1.0") return False - + print(f" ✅ Shape: {embeddings.shape}, Max norm: {max_norm:.4f}, All inside ball") - + # Check loss - loss = ckpt.get('loss', None) + loss = ckpt.get("loss", None) if loss: print(f" ✅ Loss: {loss:.6f}") - + return True - + except Exception as e: print(f" ❌ Error loading model: {e}") return False + def check_data_file(path, description): """Check data file integrity.""" try: - if path.endswith('.pkl'): - with open(path, 'rb') as f: + if path.endswith(".pkl"): + with open(path, "rb") as f: data = pickle.load(f) print(f" ✅ {len(data):,} items") - elif path.endswith('.tsv'): - df = pd.read_csv(path, sep='\t', header=None) + elif path.endswith(".tsv"): + df = pd.read_csv(path, sep="\t", header=None) print(f" ✅ {len(df):,} rows") return True except Exception as e: print(f" ❌ Error: {e}") return False + def main(): - print("="*80) + print("=" * 80) print("FINAL SANITY CHECK") - print("="*80) + print("=" * 80) print() - + all_passed = True - + # 1. Check core training scripts print("1️⃣ Core Training Scripts") all_passed &= check_file_exists("train_small.py", "Main training script") all_passed &= check_file_exists("train_hierarchical.py", "Hierarchical model") all_passed &= check_file_exists("visualize_multi_groups.py", "Visualization") print() - + # 2. Check documentation print("2️⃣ Documentation") all_passed &= check_file_exists("README.md", "Main README") @@ -94,38 +99,38 @@ def main(): all_passed &= check_file_exists("FINAL_STATUS.md", "Final status") all_passed &= check_file_exists("TRAIN_SMALL_GUIDE.md", "Training guide") print() - + # 3. Check small model (production) print("3️⃣ Small Model (Production)") - if check_file_exists("small_model_28epoch/taxonomy_model_small_best.pth", - "Best model"): - check_model("small_model_28epoch/taxonomy_model_small_best.pth", - torch.Size([92290, 10])) - - all_passed &= check_file_exists("small_model_28epoch/taxonomy_embeddings_multi_groups.png", - "Multi-group viz") - all_passed &= check_file_exists("small_model_28epoch/best_epoch_analysis_epoch28.png", - "Epoch analysis") + if check_file_exists("small_model_28epoch/taxonomy_model_small_best.pth", "Best model"): + check_model("small_model_28epoch/taxonomy_model_small_best.pth", torch.Size([92290, 10])) + + all_passed &= check_file_exists( + "small_model_28epoch/taxonomy_embeddings_multi_groups.png", "Multi-group viz" + ) + all_passed &= check_file_exists( + "small_model_28epoch/best_epoch_analysis_epoch28.png", "Epoch analysis" + ) print() - + # 4. Check animals model (reference) print("4️⃣ Animals Model (Reference)") if check_file_exists("taxonomy_model_animals_best.pth", "Animals model"): check_model("taxonomy_model_animals_best.pth", torch.Size([1055469, 10])) print() - + # 5. Check data files print("5️⃣ Data Files") if check_file_exists("data/taxonomy_edges_small_transitive.pkl", "Training data"): check_data_file("data/taxonomy_edges_small_transitive.pkl", "Training pairs") - + if check_file_exists("data/taxonomy_edges_small.mapping.tsv", "Mapping file"): check_data_file("data/taxonomy_edges_small.mapping.tsv", "TaxID mappings") - + all_passed &= check_file_exists("data/names.dmp", "NCBI names") all_passed &= check_file_exists("data/nodes.dmp", "NCBI nodes") print() - + # 6. Check no intermediate files remain print("6️⃣ Cleanup Verification") intermediate_files = [ @@ -135,25 +140,26 @@ def main(): "build_transitive_closure_full.py", "train_animals.py", ] - + cleanup_ok = True for f in intermediate_files: if os.path.exists(f): print(f" ⚠️ Intermediate file still present: {f}") cleanup_ok = False - + if cleanup_ok: print(" ✅ No intermediate files found (good!)") print() - + # Final verdict - print("="*80) + print("=" * 80) if all_passed and cleanup_ok: print("✅ ALL CHECKS PASSED - Repository is ready for commit!") else: print("❌ SOME CHECKS FAILED - Please review above") sys.exit(1) - print("="*80) + print("=" * 80) + if __name__ == "__main__": main() diff --git a/src/taxembed/visualization/__init__.py b/src/taxembed/visualization/__init__.py new file mode 100644 index 0000000..4ade23b --- /dev/null +++ b/src/taxembed/visualization/__init__.py @@ -0,0 +1,19 @@ +"""Visualization utilities for embeddings.""" + +from .umap_viz import ( + load_embeddings, + load_mapping, + load_taxonomy_tree, + visualize_multi_groups, +) +from .umap_viz import ( + main as visualize_embeddings, +) + +__all__ = [ + "load_embeddings", + "load_mapping", + "load_taxonomy_tree", + "visualize_multi_groups", + "visualize_embeddings", +] diff --git a/visualize_multi_groups.py b/src/taxembed/visualization/umap_viz.py similarity index 86% rename from visualize_multi_groups.py rename to src/taxembed/visualization/umap_viz.py index be6473f..e5defbf 100644 --- a/visualize_multi_groups.py +++ b/src/taxembed/visualization/umap_viz.py @@ -14,16 +14,18 @@ import matplotlib.pyplot as plt import numpy as np import pandas as pd -import torch import taxopy +import torch from umap import UMAP +from taxembed.data import ensure_taxdump + def load_embeddings(ckpt_path): """Load embeddings from checkpoint.""" print(f"Loading embeddings from {ckpt_path}...") ckpt = torch.load(ckpt_path, map_location="cpu") - + if "state_dict" in ckpt: sd = ckpt["state_dict"] emb = sd["lt.weight"].detach().cpu().numpy() @@ -31,7 +33,7 @@ def load_embeddings(ckpt_path): emb = ckpt["embeddings"].cpu().numpy() else: raise ValueError("Cannot find embeddings in checkpoint") - + print(f" ✓ Shape: {emb.shape}") return emb @@ -41,15 +43,15 @@ def load_mapping(map_path): if not Path(map_path).exists(): print(f"⚠️ Mapping file not found: {map_path}") return None, None - + print(f"Loading mapping from {map_path}...") df = pd.read_csv(map_path, sep="\t", dtype={"taxid": str, "idx": int}) - + # Filter out non-numeric taxids numeric_df = df[df["taxid"].str.isnumeric()] - tax2idx = dict(zip(numeric_df["taxid"], numeric_df["idx"])) - idx2tax = dict(zip(numeric_df["idx"], numeric_df["taxid"])) - + tax2idx = dict(zip(numeric_df["taxid"], numeric_df["idx"], strict=False)) + idx2tax = dict(zip(numeric_df["idx"], numeric_df["taxid"], strict=False)) + print(f" ✓ Loaded {len(tax2idx):,} mappings") return tax2idx, idx2tax @@ -58,11 +60,17 @@ def load_taxonomy_tree(valid_taxids=None, base_dir: Path = Path("data")): """Load NCBI taxonomy tree structure from dump files or TaxoPy fallback.""" names = {} nodes = {} - - # Try loading from dump files first - names_path = base_dir / "names.dmp" - nodes_path = base_dir / "nodes.dmp" - + + # Ensure dump files are extracted (will download if needed) + try: + nodes_file, names_file, _ = ensure_taxdump(base_dir) + names_path = names_file + nodes_path = nodes_file + except Exception as e: + print(f"⚠️ Could not ensure taxdump files: {e}") + names_path = base_dir / "names.dmp" + nodes_path = base_dir / "nodes.dmp" + if names_path.exists() and nodes_path.exists(): try: # Load names @@ -73,7 +81,7 @@ def load_taxonomy_tree(valid_taxids=None, base_dir: Path = Path("data")): taxid = int(parts[0]) if valid_taxids is None or taxid in valid_taxids: names[taxid] = parts[1] - + # Load nodes (parent relationships) with nodes_path.open("r") as f: for line in f: @@ -83,31 +91,31 @@ def load_taxonomy_tree(valid_taxids=None, base_dir: Path = Path("data")): parent = int(parts[1]) if valid_taxids is None or taxid in valid_taxids: nodes[taxid] = parent - - print(f"Loading taxonomy tree from dump files...") + + print("Loading taxonomy tree from dump files...") print(f" ✓ Loaded {len(nodes):,} taxonomy nodes (filtered to dataset)") return names, nodes except Exception as e: print(f"⚠️ Error reading dump files: {e}, falling back to TaxoPy...") - + # Fallback to TaxoPy try: - print(f"Loading taxonomy tree via TaxoPy (dump files not found)...") + print("Loading taxonomy tree via TaxoPy (dump files not found)...") taxdb = taxopy.TaxDb(taxdb_dir=str(base_dir)) - + # Build names dict for taxid_str, name in taxdb.taxid2name.items(): taxid = int(taxid_str) if valid_taxids is None or taxid in valid_taxids: names[taxid] = name - + # Build nodes dict (parent relationships) for taxid_str, parent_str in taxdb.taxid2parent.items(): taxid = int(taxid_str) parent = int(parent_str) if valid_taxids is None or taxid in valid_taxids: nodes[taxid] = parent - + print(f" ✓ Loaded {len(nodes):,} taxonomy nodes via TaxoPy (filtered to dataset)") return names, nodes except Exception as e: @@ -140,7 +148,7 @@ def get_nodes_at_depth(root_taxid, parent_children, depth): """Get all nodes at a specific depth from root (0=children, 1=grandchildren, etc.).""" if depth == 0: return parent_children.get(root_taxid, []) - + current_level = [root_taxid] for _ in range(depth): next_level = [] @@ -167,7 +175,7 @@ def visualize_multi_groups( loss=None, ): """Create UMAP visualization with multiple highlighted groups. - + Args: child_coloring: Root TaxID to color by children coloring_depth: Depth level for coloring (0=children, 1=grandchildren, 2=great-grandchildren, etc.) @@ -175,28 +183,34 @@ def visualize_multi_groups( epoch: Training epoch for title loss: Training loss for title """ - + # Find members of each group print("\nFinding taxonomic groups...") group_members = {} if child_coloring is not None: parent_children = build_parent_children(nodes) root_taxid = child_coloring - + # Get nodes at the specified depth depth_nodes = get_nodes_at_depth(root_taxid, parent_children, coloring_depth) - + if not depth_nodes: - depth_label = ["children", "grandchildren", "great-grandchildren"][min(coloring_depth, 2)] + depth_label = ["children", "grandchildren", "great-grandchildren"][ + min(coloring_depth, 2) + ] if coloring_depth > 2: depth_label = f"{coloring_depth}-level descendants" - print(f" ⚠️ No {depth_label} found at depth {coloring_depth} under root TaxID {root_taxid}") + print( + f" ⚠️ No {depth_label} found at depth {coloring_depth} under root TaxID {root_taxid}" + ) else: - depth_label = ["children", "grandchildren", "great-grandchildren"][min(coloring_depth, 2)] + depth_label = ["children", "grandchildren", "great-grandchildren"][ + min(coloring_depth, 2) + ] if coloring_depth > 2: depth_label = f"{coloring_depth}-level descendants" print(f" Coloring by {depth_label} (depth {coloring_depth})...") - + for node in depth_nodes: # Collect all descendants of this node for coloring indices = collect_descendants(node, parent_children, tax2idx) @@ -221,11 +235,11 @@ def visualize_multi_groups( else: print(f" ⚠️ {group_name}: root TaxID {root_taxid} not found in taxonomy tree") group_members[group_name] = set() - + # Assign colors to all indices n_total = emb.shape[0] idx_to_group = {} - + for idx in range(n_total): assigned = False for group_name, members in group_members.items(): @@ -235,15 +249,15 @@ def visualize_multi_groups( break if not assigned: idx_to_group[idx] = "Other" - + # Sample points print(f"\nSampling {sample_size:,} points from {n_total:,} total...") all_indices = list(range(n_total)) - + # Stratified sampling: ensure each group is represented sampled_indices = [] samples_per_group = {} - + label_order = list(group_members.keys()) if not label_order: label_order = ["Other"] @@ -258,46 +272,50 @@ def visualize_multi_groups( group_indices = [i for i in all_indices if idx_to_group[i] == group_name] if group_indices: # Sample proportionally - n_sample = min(len(group_indices), max(1, int(len(group_indices) / n_total * sample_size))) + n_sample = min( + len(group_indices), max(1, int(len(group_indices) / n_total * sample_size)) + ) sampled = np.random.choice(group_indices, n_sample, replace=False) sampled_indices.extend(sampled) samples_per_group[group_name] = len(sampled) - + # If we haven't reached sample_size, add more from "Other" if len(sampled_indices) < sample_size: other_indices = [i for i in all_indices if i not in sampled_indices] - additional = np.random.choice(other_indices, - min(len(other_indices), sample_size - len(sampled_indices)), - replace=False) + additional = np.random.choice( + other_indices, + min(len(other_indices), sample_size - len(sampled_indices)), + replace=False, + ) sampled_indices.extend(additional) - + sampled_indices = np.array(sampled_indices) - + print(f" ✓ Sampled {len(sampled_indices):,} points") for group_name in label_order + ["Other"]: if group_name in samples_per_group: print(f" - {group_name}: {samples_per_group[group_name]:,}") - + # Extract embeddings sample_emb = emb[sampled_indices] - + # Run UMAP print(f"\nRunning UMAP on {len(sampled_indices):,} points...") umap_model = UMAP(n_components=2, random_state=42, n_neighbors=15, min_dist=0.1) projection = umap_model.fit_transform(sample_emb) - + # Plot print("Creating visualization...") fig, ax = plt.subplots(figsize=(18, 14)) - + # Plot each group for group_name in ["Other"] + label_order: mask = np.array([idx_to_group[idx] == group_name for idx in sampled_indices]) n_points = mask.sum() - + if n_points > 0: color = color_lookup.get(group_name, "#cccccc") - + # Other group: smaller, more transparent if group_name == "Other": ax.scatter( @@ -308,7 +326,7 @@ def visualize_multi_groups( alpha=0.2, label=f"{group_name} (n={n_points:,})", edgecolors="none", - zorder=1 + zorder=1, ) else: # Highlighted groups: larger, more visible @@ -321,37 +339,37 @@ def visualize_multi_groups( label=f"{group_name} (n={n_points:,})", edgecolors="black", linewidth=0.5, - zorder=2 + zorder=2, ) - + ax.set_xlabel("UMAP 1", fontsize=16) ax.set_ylabel("UMAP 2", fontsize=16) - + # Build title title_parts = [] if clade_name: title_parts.append(f"TaxEmbed: {clade_name}") else: title_parts.append("TaxEmbed") - + title_parts.append(f"Children Level {coloring_depth}") - + if epoch is not None: title_parts.append(f"epochs {epoch}") - + if loss is not None: title_parts.append(f"Loss {loss:.6f}") - + title = ", ".join(title_parts) ax.set_title(title, fontsize=20, fontweight="bold") ax.legend(loc="best", fontsize=13, framealpha=0.95) ax.grid(True, alpha=0.3) - + plt.tight_layout() - + if output_file is None: output_file = "taxonomy_embeddings_multi_groups.png" - + plt.savefig(output_file, dpi=200, bbox_inches="tight") print(f"\n✅ Saved: {output_file}") @@ -365,46 +383,54 @@ def main(): parser.add_argument("--names", help="Path to names.dmp (default: data/names.dmp)") parser.add_argument("--nodes", help="Path to nodes.dmp (default: data/nodes.dmp)") parser.add_argument("--root-taxid", type=int, help="Root TaxID for child-level coloring") - parser.add_argument("--children", type=int, default=0, - help="Depth level for coloring (0=children, 1=grandchildren, 2=great-grandchildren, etc.)") + parser.add_argument( + "--children", + type=int, + default=0, + help="Depth level for coloring (0=children, 1=grandchildren, 2=great-grandchildren, etc.)", + ) parser.add_argument("--clade-name", help="Name of the clade for title") parser.add_argument("--epoch", type=int, help="Training epoch for title") parser.add_argument("--loss", type=float, help="Training loss for title") - + args = parser.parse_args() - + # Load embeddings emb = load_embeddings(args.checkpoint) - + # Auto-detect mapping file if args.mapping is None: # Try to find mapping file - for candidate in ["data/taxonomy_edges_small.mapping.tsv", - "data/taxonomy_edges.mapping.tsv"]: + for candidate in [ + "data/taxonomy_edges_small.mapping.tsv", + "data/taxonomy_edges.mapping.tsv", + ]: if Path(candidate).exists(): args.mapping = candidate break - + if args.mapping is None: print("❌ Could not find mapping file. Please specify with --mapping") sys.exit(1) - + # Load mapping tax2idx, idx2tax = load_mapping(args.mapping) if tax2idx is None: sys.exit(1) - + # Convert to set of numeric taxids - valid_taxids = set(int(t) for t in tax2idx.keys()) + valid_taxids = {int(t) for t in tax2idx.keys()} print(f"Dataset contains {len(valid_taxids):,} unique organisms") - - # Load taxonomy tree - find data directory relative to script location - script_dir = Path(__file__).resolve().parent - default_data_dir = script_dir / "data" - + + # Load taxonomy tree - find data directory at project root + # Script is at: src/taxembed/visualization/umap_viz.py + # Project root is 3 levels up + project_root = Path(__file__).resolve().parents[3] + default_data_dir = project_root / "data" + if args.names or args.nodes: names_path = Path(args.names).resolve() if args.names else default_data_dir / "names.dmp" - nodes_path = Path(args.nodes).resolve() if args.nodes else names_path.parent / "nodes.dmp" + Path(args.nodes).resolve() if args.nodes else names_path.parent / "nodes.dmp" base_dir = names_path.parent names, nodes = load_taxonomy_tree(valid_taxids, base_dir=base_dir) else: @@ -413,12 +439,12 @@ def main(): if names is None or nodes is None: print("❌ Could not load taxonomy tree") sys.exit(1) - + # Prepare child coloring if root taxid provided child_coloring = None if args.root_taxid is not None: child_coloring = args.root_taxid - + # Visualize visualize_multi_groups( emb, diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..4f1fe1e --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,58 @@ +"""Pytest configuration and fixtures.""" + +import pytest + + +@pytest.fixture +def sample_training_data(): + """Create sample training data for tests.""" + return [ + { + "ancestor_idx": 0, + "descendant_idx": 1, + "ancestor_depth": 0, + "descendant_depth": 1, + "depth_diff": 1, + }, + { + "ancestor_idx": 0, + "descendant_idx": 2, + "ancestor_depth": 0, + "descendant_depth": 2, + "depth_diff": 2, + }, + { + "ancestor_idx": 1, + "descendant_idx": 2, + "ancestor_depth": 1, + "descendant_depth": 2, + "depth_diff": 1, + }, + ] + + +@pytest.fixture +def sample_depth_map(): + """Create sample depth mapping.""" + return { + 0: 0, # Root + 1: 1, # Level 1 + 2: 2, # Level 2 + 3: 2, # Level 2 + } + + +@pytest.fixture +def simple_model(): + """Create a simple Poincaré model for testing.""" + from taxembed.models import HierarchicalPoincareEmbedding + + return HierarchicalPoincareEmbedding(n_nodes=10, dim=5, max_depth=3) + + +@pytest.fixture +def data_dir(tmp_path): + """Create temporary data directory.""" + data_dir = tmp_path / "data" + data_dir.mkdir() + return data_dir diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..ab4cfbc --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,43 @@ +"""Tests for data processing utilities.""" + +from taxembed.data import parse_names_dmp, parse_nodes_dmp + + +class TestDataParsing: + """Test data parsing functions.""" + + def test_parse_nodes_dmp(self, tmp_path): + """Test parsing nodes.dmp file.""" + # Create a temporary nodes.dmp file + nodes_file = tmp_path / "nodes.dmp" + nodes_file.write_text( + "1\t|\t1\t|\tno rank\n" # Root (self-loop, should be skipped) + "2\t|\t1\t|\tspecies\n" + "3\t|\t1\t|\tspecies\n" + "4\t|\t2\t|\tsubspecies\n" + ) + + edges = parse_nodes_dmp(nodes_file) + + # Should have 3 edges (root self-loop excluded) + assert len(edges) == 3 + + # Check structure + assert all("id1" in e and "id2" in e for e in edges) + + def test_parse_names_dmp(self, tmp_path): + """Test parsing names.dmp file.""" + # Create a temporary names.dmp file + names_file = tmp_path / "names.dmp" + names_file.write_text( + "1\t|\troot\t|\t\t|\tscientific name\t|\n" + "2\t|\tBacteria\t|\t\t|\tscientific name\t|\n" + "3\t|\tArchaea\t|\t\t|\tscientific name\t|\n" + ) + + names_map = parse_names_dmp(names_file) + + assert len(names_map) == 3 + assert names_map[1] == "root" + assert names_map[2] == "Bacteria" + assert names_map[3] == "Archaea" diff --git a/tests/test_example.py b/tests/test_example.py deleted file mode 100644 index ddc52a8..0000000 --- a/tests/test_example.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Example test module.""" - - -def test_placeholder(): - """Placeholder test.""" - assert True diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..6cf3e9f --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,110 @@ +"""Tests for Poincaré embedding models.""" + +import torch + +from taxembed.models import HierarchicalPoincareEmbedding, MetricsTracker + + +class TestHierarchicalPoincareEmbedding: + """Test the Poincaré embedding model.""" + + def test_initialization(self): + """Test model initialization.""" + model = HierarchicalPoincareEmbedding(n_nodes=100, dim=10, max_depth=5) + assert model.n_nodes == 100 + assert model.dim == 10 + assert model.max_depth == 5 + assert model.embeddings.weight.shape == (100, 10) + + def test_depth_initialization(self, sample_depth_map): + """Test depth-aware initialization.""" + model = HierarchicalPoincareEmbedding( + n_nodes=4, dim=5, max_depth=2, init_depth_data=sample_depth_map + ) + + # Check that norms increase with depth + norms = model.embeddings.weight.norm(dim=1) + assert norms[0] < norms[1] < norms[2] # Depth 0 < 1 < 2 + + def test_forward(self, simple_model): + """Test forward pass.""" + indices = torch.tensor([0, 1, 2]) + embeddings = simple_model(indices) + assert embeddings.shape == (3, 5) + + def test_poincare_distance(self, simple_model): + """Test Poincaré distance computation.""" + u = torch.randn(10, 5) * 0.5 # Keep inside ball + v = torch.randn(10, 5) * 0.5 + + distances = simple_model.poincare_distance(u, v) + assert distances.shape == (10,) + assert torch.all(distances >= 0) # Distances are non-negative + + def test_project_to_ball(self, simple_model): + """Test ball projection.""" + # Set some embeddings outside the ball + with torch.no_grad(): + simple_model.embeddings.weight[0] = torch.ones(5) * 1.5 + + # Project back + simple_model.project_to_ball() + + # Check all embeddings are inside ball + norms = simple_model.embeddings.weight.norm(dim=1) + assert torch.all(norms < 1.0) + + def test_selective_projection(self, simple_model): + """Test selective projection of specific indices.""" + # Set one embedding outside + with torch.no_grad(): + simple_model.embeddings.weight[5] = torch.ones(5) * 1.5 + + # Project only that index + simple_model.project_to_ball(indices=torch.tensor([5])) + + # Check it's inside + norm = simple_model.embeddings.weight[5].norm() + assert norm < 1.0 + + +class TestMetricsTracker: + """Test the metrics tracker.""" + + def test_initialization(self): + """Test tracker initialization.""" + tracker = MetricsTracker() + assert len(tracker.history) == 0 + assert tracker.best_loss == float("inf") + assert tracker.best_epoch == 0 + + def test_update(self): + """Test metric updates.""" + tracker = MetricsTracker() + metrics = {"loss": 0.5, "reg_loss": 0.1} + + tracker.update(1, metrics) + assert len(tracker.history) == 1 + assert tracker.best_loss == 0.5 + assert tracker.best_epoch == 1 + + def test_best_loss_tracking(self): + """Test that best loss is tracked correctly.""" + tracker = MetricsTracker() + + tracker.update(1, {"loss": 0.5}) + tracker.update(2, {"loss": 0.3}) # Better + tracker.update(3, {"loss": 0.4}) # Worse + + assert tracker.best_loss == 0.3 + assert tracker.best_epoch == 2 + + def test_get_previous(self): + """Test getting previous metrics.""" + tracker = MetricsTracker() + + tracker.update(1, {"loss": 0.5}) + assert tracker.get_previous("loss") is None # No previous + + tracker.update(2, {"loss": 0.3}) + assert tracker.get_previous("loss") == 0.5 diff --git a/tests/test_training.py b/tests/test_training.py new file mode 100644 index 0000000..1b3cc60 --- /dev/null +++ b/tests/test_training.py @@ -0,0 +1,90 @@ +"""Tests for training utilities.""" + +import torch + +from taxembed.training import ( + HierarchicalDataLoader, + radial_regularizer, + ranking_loss_with_margin, +) + + +class TestHierarchicalDataLoader: + """Test the hierarchical data loader.""" + + def test_initialization(self, sample_training_data): + """Test data loader initialization.""" + loader = HierarchicalDataLoader( + training_data=sample_training_data, n_nodes=10, batch_size=2, n_negatives=5 + ) + assert loader.batch_size == 2 + assert loader.n_negatives == 5 + + def test_length(self, sample_training_data): + """Test data loader length.""" + loader = HierarchicalDataLoader( + training_data=sample_training_data, n_nodes=10, batch_size=2, n_negatives=5 + ) + # 3 items with batch_size 2 = 1 batch (3 // 2) + assert len(loader) == 1 + + def test_iteration(self, sample_training_data): + """Test iterating through data loader.""" + loader = HierarchicalDataLoader( + training_data=sample_training_data, n_nodes=10, batch_size=2, n_negatives=5 + ) + + for ancestors, descendants, negatives, depths in loader: + assert ancestors.shape[0] <= 2 # Batch size + assert descendants.shape[0] <= 2 + assert negatives.shape == (ancestors.shape[0], 5) # n_negatives + assert depths.shape[0] <= 2 + break # Just test one batch + + +class TestLossFunctions: + """Test loss functions.""" + + def test_ranking_loss_with_margin(self, simple_model): + """Test ranking loss computation.""" + batch_size = 4 + n_negatives = 3 + + ancestors = torch.randint(0, 10, (batch_size,)) + descendants = torch.randint(0, 10, (batch_size,)) + negatives = torch.randint(0, 10, (batch_size, n_negatives)) + depths = torch.rand(batch_size) + + loss = ranking_loss_with_margin( + simple_model, + ancestors, + descendants, + negatives, + depths, + margin=0.1, + depth_weight=True, + ) + + assert isinstance(loss, torch.Tensor) + assert loss.ndim == 0 # Scalar + assert loss >= 0 # Loss should be non-negative + + def test_radial_regularizer(self, simple_model, sample_depth_map): + """Test radial regularization.""" + idx_tensor = torch.tensor([0, 1, 2]) + target_radii = torch.tensor([0.1, 0.5, 0.9]) + + reg_loss = radial_regularizer(simple_model, idx_tensor, target_radii, lambda_reg=0.1) + + assert isinstance(reg_loss, torch.Tensor) + assert reg_loss.ndim == 0 # Scalar + assert reg_loss >= 0 + + def test_radial_regularizer_empty(self, simple_model): + """Test radial regularizer with empty indices.""" + idx_tensor = torch.tensor([], dtype=torch.long) + target_radii = torch.tensor([]) + + reg_loss = radial_regularizer(simple_model, idx_tensor, target_radii, lambda_reg=0.1) + + assert reg_loss == 0.0 diff --git a/tests/test_validation.py b/tests/test_validation.py new file mode 100644 index 0000000..08166e6 --- /dev/null +++ b/tests/test_validation.py @@ -0,0 +1,29 @@ +"""Tests for validation utilities.""" + +import torch + +from taxembed.models import HierarchicalPoincareEmbedding + + +class TestValidation: + """Test validation functions.""" + + def test_ball_constraint_validation(self, simple_model): + """Test that ball constraints are maintained.""" + # All embeddings should be inside unit ball after projection + simple_model.project_to_ball() + + norms = simple_model.embeddings.weight.norm(dim=1) + assert torch.all(norms < 1.0), "Some embeddings are outside the unit ball" + + def test_depth_ordering(self, sample_depth_map): + """Test that deeper nodes have larger norms.""" + model = HierarchicalPoincareEmbedding( + n_nodes=4, dim=5, max_depth=2, init_depth_data=sample_depth_map + ) + + norms = model.embeddings.weight.norm(dim=1) + + # Root should have smallest norm + assert norms[0] < norms[1] + assert norms[0] < norms[2] diff --git a/train_hierarchical.py b/train_hierarchical.py deleted file mode 100644 index 4940d27..0000000 --- a/train_hierarchical.py +++ /dev/null @@ -1,583 +0,0 @@ -#!/usr/bin/env python3 -""" -Phase 2: Hierarchical Training for Poincaré Embeddings - -Key improvements: -1. Train on transitive closure (ALL ancestor-descendant pairs) -2. Depth-aware initialization (deeper nodes near boundary) -3. Radial regularizer (enforce depth → radius mapping) -4. Hard negative sampling (cousins at same depth) -5. Depth weighting (deeper pairs matter more) -6. Proper hyperbolic distance and loss -""" - -import torch -import torch.nn as nn -import torch.optim as optim -import numpy as np -import pickle -from collections import defaultdict, deque -from tqdm import tqdm -import argparse -import os - - -class HierarchicalPoincareEmbedding(nn.Module): - """Poincaré embeddings with hierarchical structure.""" - - def __init__(self, n_nodes, dim=10, max_depth=38, init_depth_data=None): - super().__init__() - self.n_nodes = n_nodes - self.dim = dim - self.max_depth = max_depth - - # Embeddings (initialize later with depth info) - self.embeddings = nn.Embedding(n_nodes, dim) - - # Initialize based on depth if available - if init_depth_data is not None: - self._initialize_by_depth(init_depth_data) - else: - # Default: uniform small initialization - nn.init.uniform_(self.embeddings.weight, -0.001, 0.001) - - def _initialize_by_depth(self, depth_data): - """ - Initialize embeddings based on taxonomic depth. - - Deeper nodes → larger radius (closer to boundary). - This encodes hierarchy from the start! - """ - print("Initializing embeddings by depth...") - - # Map: idx → depth - idx_to_depth = depth_data - - with torch.no_grad(): - for idx in range(self.n_nodes): - depth = idx_to_depth.get(idx, 0) - - # Radius increases with depth - # Root (depth 0): r ≈ 0.1 - # Max depth: r ≈ 0.95 - target_radius = 0.1 + (depth / self.max_depth) * 0.85 - - # Random direction on sphere - vec = torch.randn(self.dim) - vec = vec / vec.norm() - - # Scale to target radius - self.embeddings.weight[idx] = vec * target_radius - - norms = self.embeddings.weight.norm(dim=1) - print(f" ✓ Initialized: norm range [{norms.min():.3f}, {norms.max():.3f}]") - - def forward(self, indices): - """Get embeddings for indices.""" - return self.embeddings(indices) - - def poincare_distance(self, u, v, eps=1e-5): - """ - Compute Poincaré distance between embeddings. - - d(u,v) = arcosh(1 + 2||u-v||²/((1-||u||²)(1-||v||²))) - """ - # Compute squared norms - u_norm_sq = (u ** 2).sum(dim=-1) - v_norm_sq = (v ** 2).sum(dim=-1) - - # Clamp to stay inside ball - u_norm_sq = torch.clamp(u_norm_sq, 0, 1 - eps) - v_norm_sq = torch.clamp(v_norm_sq, 0, 1 - eps) - - # Squared Euclidean distance - diff_norm_sq = ((u - v) ** 2).sum(dim=-1) - - # Poincaré distance - numerator = 2 * diff_norm_sq - denominator = (1 - u_norm_sq) * (1 - v_norm_sq) - - dist = torch.acosh(1 + numerator / (denominator + eps) + eps) - - return dist - - def project_to_ball(self, indices=None, max_norm=0.999): - """ - Project embeddings back into Poincaré ball with HARD constraint. - - Args: - indices: If provided, only project these indices (more efficient). - If None, project all embeddings. - max_norm: Maximum allowed norm (default 0.999, essentially at boundary) - """ - with torch.no_grad(): - if indices is not None: - # Only project updated embeddings - embs = self.embeddings.weight[indices] - norms = embs.norm(dim=1, keepdim=True) - # Hard projection: if norm >= max_norm, scale it down - # Use where to only scale embeddings that need it - needs_projection = norms >= max_norm - scale = torch.where( - needs_projection, - max_norm / (norms + 1e-8), - torch.ones_like(norms) - ) - self.embeddings.weight[indices] = embs * scale - else: - # Project all embeddings - norms = self.embeddings.weight.norm(dim=1, keepdim=True) - needs_projection = norms >= max_norm - scale = torch.where( - needs_projection, - max_norm / (norms + 1e-8), - torch.ones_like(norms) - ) - self.embeddings.weight.mul_(scale) - - -class HierarchicalDataLoader: - """ - Data loader with depth-aware sampling and hard negatives. - """ - - def __init__(self, training_data, n_nodes, batch_size=32, - n_negatives=50, depth_stratify=True): - self.training_data = training_data # List of dicts with metadata - self.n_nodes = n_nodes - self.batch_size = batch_size - self.n_negatives = n_negatives - self.depth_stratify = depth_stratify - - # Build index by depth for stratified sampling - if depth_stratify: - self.depth_buckets = defaultdict(list) - for i, item in enumerate(training_data): - depth_diff = item['depth_diff'] - self.depth_buckets[depth_diff].append(i) - print(f" ✓ Created {len(self.depth_buckets)} depth buckets for sampling") - - # Build node → siblings map for hard negatives - self._build_sibling_map() - - def _build_sibling_map(self): - """Build map: node → nodes at same depth (for hard negatives).""" - print(" Building sibling map for hard negatives...") - - depth_to_nodes = defaultdict(set) - for item in self.training_data: - depth_to_nodes[item['descendant_depth']].add(item['descendant_idx']) - - self.sibling_map = {} - for depth, nodes in depth_to_nodes.items(): - nodes_list = list(nodes) - for node in nodes_list: - # Siblings = other nodes at same depth - self.sibling_map[node] = [n for n in nodes_list if n != node] - - print(f" ✓ Built sibling map for hard negatives") - - def __len__(self): - return len(self.training_data) // self.batch_size - - def __iter__(self): - """Iterate over batches with depth-aware sampling.""" - indices = list(range(len(self.training_data))) - - if self.depth_stratify: - # Stratified sampling: mix shallow and deep pairs - np.random.shuffle(indices) - else: - # Random sampling - np.random.shuffle(indices) - - for i in range(0, len(indices), self.batch_size): - batch_indices = indices[i:i + self.batch_size] - batch = [self.training_data[idx] for idx in batch_indices] - - # Extract data (vectorized for speed) - batch_size = len(batch) - ancestors = np.zeros(batch_size, dtype=np.int64) - descendants = np.zeros(batch_size, dtype=np.int64) - depths = np.zeros(batch_size, dtype=np.float32) - - for j, item in enumerate(batch): - ancestors[j] = item['ancestor_idx'] - descendants[j] = item['descendant_idx'] - depths[j] = item['depth_diff'] - - # Convert to tensors once - ancestors = torch.from_numpy(ancestors) - descendants = torch.from_numpy(descendants) - depths = torch.from_numpy(depths) - - # Sample hard negatives (cousins at same depth) - negatives = np.zeros((batch_size, self.n_negatives), dtype=np.int64) - for j, item in enumerate(batch): - desc_idx = item['descendant_idx'] - - # Get siblings (nodes at same depth) - siblings = self.sibling_map.get(desc_idx, []) - - if len(siblings) >= self.n_negatives: - # Sample from siblings (hard negatives) - negatives[j] = np.random.choice(siblings, self.n_negatives, replace=False) - else: - # Mix siblings + random negatives - n_sibling = len(siblings) - n_random = self.n_negatives - n_sibling - if n_sibling > 0: - negatives[j, :n_sibling] = siblings - negatives[j, n_sibling:] = np.random.choice(self.n_nodes, n_random, replace=False) - else: - negatives[j] = np.random.choice(self.n_nodes, self.n_negatives, replace=False) - - negatives = torch.from_numpy(negatives) - - yield ancestors, descendants, negatives, depths - - -def ranking_loss_with_margin(model, ancestors, descendants, negatives, - depths, margin=0.1, depth_weight=True): - """ - Ranking loss with margin and optional depth weighting. - - Loss encourages: - - d(ancestor, descendant) < d(ancestor, negative) + margin - - Deeper pairs get higher weight (they're more informative) - """ - # Get embeddings - anc_emb = model(ancestors) # (batch, dim) - desc_emb = model(descendants) # (batch, dim) - neg_emb = model(negatives) # (batch, n_neg, dim) - - # Positive distances (ancestor → descendant) - pos_dist = model.poincare_distance(anc_emb, desc_emb) # (batch,) - - # Negative distances (ancestor → each negative) - # Expand anc_emb to match negatives shape - anc_emb_expanded = anc_emb.unsqueeze(1).expand_as(neg_emb) # (batch, n_neg, dim) - neg_dist = model.poincare_distance(anc_emb_expanded, neg_emb) # (batch, n_neg) - - # Margin ranking loss: max(0, pos_dist - neg_dist + margin) - losses = torch.relu(pos_dist.unsqueeze(1) - neg_dist + margin) # (batch, n_neg) - loss = losses.mean(dim=1) # Average over negatives: (batch,) - - # Depth weighting: deeper pairs are more important - if depth_weight: - # Weight = sqrt(depth) to emphasize deep pairs without over-weighting - weights = torch.sqrt(depths + 1) # +1 to avoid zero weight - weights = weights / weights.mean() # Normalize - loss = loss * weights - - return loss.mean() - - -def radial_regularizer(model, idx_to_depth_tensor, target_radii_tensor, lambda_reg=0.01): - """ - Vectorized radial regularizer to keep nodes at expected radius based on depth. - - Encourages: ||embedding|| ≈ f(depth) - where f(depth) = 0.1 + (depth/max_depth) * 0.85 - - Args: - idx_to_depth_tensor: Tensor of indices to regularize - target_radii_tensor: Tensor of target radii for each index - """ - if len(idx_to_depth_tensor) == 0: - return torch.tensor(0.0, device=model.embeddings.weight.device) - - # Get embeddings for these indices - embs = model.embeddings.weight[idx_to_depth_tensor] # (n, dim) - - # Compute actual radii - actual_radii = embs.norm(dim=1) # (n,) - - # L2 penalty - reg_loss = ((actual_radii - target_radii_tensor) ** 2).mean() - - return lambda_reg * reg_loss - - -def train_hierarchical(model, dataloader, optimizer, n_epochs, - idx_to_depth, max_depth, device, - margin=0.2, lambda_reg=0.01, early_stopping_patience=3, - checkpoint_base=None): - """Train with hierarchical constraints and early stopping.""" - - model.to(device) - model.train() - - print(f"\nStarting hierarchical training...") - print(f" Margin: {margin}") - print(f" Radial regularization: λ={lambda_reg}") - print(f" Early stopping patience: {early_stopping_patience} epochs") - print(f" Device: {device}") - print() - - # Precompute regularizer tensors (once, not every batch!) - print("Precomputing radial regularization tensors...") - reg_indices = [] - reg_target_radii = [] - for idx, depth in idx_to_depth.items(): - if idx < model.n_nodes: - reg_indices.append(idx) - # Match initialization: 0.1 + depth/max * 0.85 - target_radius = 0.1 + (depth / max_depth) * 0.85 - reg_target_radii.append(target_radius) - - reg_indices_tensor = torch.LongTensor(reg_indices).to(device) - reg_target_radii_tensor = torch.FloatTensor(reg_target_radii).to(device) - print(f" ✓ Will regularize {len(reg_indices):,} nodes") - - # Early stopping tracking - best_loss = float('inf') - epochs_without_improvement = 0 - - # Checkpoint management (keep only last 5 epoch checkpoints) - checkpoint_queue = deque(maxlen=5) - - for epoch in range(n_epochs): - epoch_loss = 0 - epoch_reg_loss = 0 - n_batches = 0 - - pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{n_epochs}") - - for ancestors, descendants, negatives, depths in pbar: - ancestors = ancestors.to(device) - descendants = descendants.to(device) - negatives = negatives.to(device) - depths = depths.to(device) - - optimizer.zero_grad() - - # Ranking loss - loss = ranking_loss_with_margin( - model, ancestors, descendants, negatives, - depths, margin=margin, depth_weight=True - ) - - # Radial regularizer (using precomputed tensors) - reg_loss = radial_regularizer(model, reg_indices_tensor, reg_target_radii_tensor, lambda_reg) - - # Total loss - total_loss = loss + reg_loss - - # Backward - total_loss.backward() - - # Clip gradients to prevent exploding gradients - torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) - - optimizer.step() - - # Project back to ball (only modified embeddings for efficiency) - # Collect all unique indices that were updated - updated_indices = torch.cat([ancestors, descendants, negatives.flatten()]) - updated_indices = torch.unique(updated_indices) - model.project_to_ball(updated_indices) - - # Periodic full projection to catch any stragglers - # Every 500 batches, project ALL embeddings - if n_batches % 500 == 0: - model.project_to_ball(indices=None) # Project all - - # Track - epoch_loss += loss.item() - epoch_reg_loss += reg_loss.item() - n_batches += 1 - - # Update progress - pbar.set_postfix({ - 'loss': f'{loss.item():.4f}', - 'reg': f'{reg_loss.item():.4f}' - }) - - # Epoch summary - avg_loss = epoch_loss / n_batches - avg_reg = epoch_reg_loss / n_batches - - # FINAL projection: enforce ALL embeddings are inside ball at epoch end - model.project_to_ball(indices=None) - - print(f"Epoch {epoch+1:3d}: Loss={avg_loss:.4f}, Reg={avg_reg:.4f}") - - # Check norms (after final projection) - norms = model.embeddings.weight.norm(dim=1) - outside_count = (norms >= 1.0).sum().item() - print(f" Norms: min={norms.min():.4f}, mean={norms.mean():.4f}, max={norms.max():.4f}") - if outside_count > 0: - print(f" ⚠️ {outside_count} embeddings still outside ball (should be 0!)") - - # Save checkpoint every epoch - if checkpoint_base: - checkpoint_path = checkpoint_base.replace('.pth', f'_epoch{epoch+1}.pth') - torch.save({ - 'state_dict': {'lt.weight': model.embeddings.weight}, - 'embeddings': model.embeddings.weight, - 'epoch': epoch + 1, - 'loss': avg_loss, - 'reg_loss': avg_reg, - 'best_loss': best_loss, - 'epochs_without_improvement': epochs_without_improvement, - }, checkpoint_path) - - # Manage checkpoint queue (keep only last 5) - if len(checkpoint_queue) >= checkpoint_queue.maxlen: - old_checkpoint = checkpoint_queue[0] - if os.path.exists(old_checkpoint): - os.remove(old_checkpoint) - print(f" 🗑️ Deleted old: {os.path.basename(old_checkpoint)}") - - checkpoint_queue.append(checkpoint_path) - print(f" 💾 Saved: {os.path.basename(checkpoint_path)} (keeping last {len(checkpoint_queue)})") - - # Early stopping check - if avg_loss < best_loss: - improvement = best_loss - avg_loss - print(f" ✓ Loss improved by {improvement:.6f}") - best_loss = avg_loss - epochs_without_improvement = 0 - - # Save best model - if checkpoint_base: - best_checkpoint = checkpoint_base.replace('.pth', '_best.pth') - torch.save({ - 'state_dict': {'lt.weight': model.embeddings.weight}, - 'embeddings': model.embeddings.weight, - 'epoch': epoch + 1, - 'loss': avg_loss, - 'reg_loss': avg_reg, - }, best_checkpoint) - print(f" 💾 Best model saved: {best_checkpoint}") - else: - epochs_without_improvement += 1 - print(f" ✗ No improvement ({epochs_without_improvement}/{early_stopping_patience})") - - if epochs_without_improvement >= early_stopping_patience: - print(f"\n🛑 Early stopping triggered after {epoch+1} epochs") - print(f" Best loss: {best_loss:.6f}") - break - - -def main(): - parser = argparse.ArgumentParser(description='Hierarchical Poincaré Training') - parser.add_argument('--data', default='data/taxonomy_edges_small_transitive.pkl', - help='Training data (pickle file with metadata)') - parser.add_argument('--checkpoint', default='taxonomy_model_hierarchical.pth', - help='Output checkpoint path') - parser.add_argument('--dim', type=int, default=10, - help='Embedding dimension') - parser.add_argument('--epochs', type=int, default=10000, - help='Maximum number of epochs (early stopping will trigger)') - parser.add_argument('--early-stopping', type=int, default=3, - help='Early stopping patience (stop if no improvement for N epochs)') - parser.add_argument('--batch-size', type=int, default=64, - help='Batch size') - parser.add_argument('--n-negatives', type=int, default=50, - help='Number of negative samples') - parser.add_argument('--lr', type=float, default=0.005, - help='Learning rate (reduced to prevent escaping ball)') - parser.add_argument('--margin', type=float, default=0.2, - help='Ranking loss margin') - parser.add_argument('--lambda-reg', type=float, default=0.1, - help='Radial regularization weight (increased from 0.01 to keep embeddings in ball)') - parser.add_argument('--gpu', type=int, default=-1, - help='GPU device (-1 for CPU)') - - args = parser.parse_args() - - print("="*80) - print("HIERARCHICAL POINCARÉ TRAINING") - print("="*80) - print() - - # Device (force CPU for stability on macOS - MPS can hang with custom ops) - if args.gpu >= 0 and torch.cuda.is_available(): - device = torch.device(f'cuda:{args.gpu}') - print(f"Using GPU: cuda:{args.gpu}") - else: - device = torch.device('cpu') - print(f"Using CPU (recommended for macOS)") - - # Note: MPS disabled because it can hang with hyperbolic distance operations - - # Load training data - print(f"Loading training data from {args.data}...") - with open(args.data, 'rb') as f: - training_data = pickle.load(f) - print(f" ✓ Loaded {len(training_data):,} training pairs") - - # Get number of nodes and max depth - n_nodes = max(max(item['ancestor_idx'], item['descendant_idx']) - for item in training_data) + 1 - max_depth = max(item['descendant_depth'] for item in training_data) - - print(f" Nodes: {n_nodes:,}") - print(f" Max depth: {max_depth}") - - # Build idx → depth mapping for initialization - idx_to_depth = {} - for item in training_data: - idx_to_depth[item['descendant_idx']] = item['descendant_depth'] - if item['ancestor_idx'] not in idx_to_depth: - idx_to_depth[item['ancestor_idx']] = item['ancestor_depth'] - - # Create model with depth-aware initialization - print("\nCreating model...") - model = HierarchicalPoincareEmbedding( - n_nodes=n_nodes, - dim=args.dim, - max_depth=max_depth, - init_depth_data=idx_to_depth - ) - - # Create dataloader - print("\nCreating dataloader...") - dataloader = HierarchicalDataLoader( - training_data=training_data, - n_nodes=n_nodes, - batch_size=args.batch_size, - n_negatives=args.n_negatives, - depth_stratify=True - ) - - # Optimizer - optimizer = optim.Adam(model.parameters(), lr=args.lr) - - # Train - train_hierarchical( - model=model, - dataloader=dataloader, - optimizer=optimizer, - n_epochs=args.epochs, - idx_to_depth=idx_to_depth, - max_depth=max_depth, - device=device, - margin=args.margin, - lambda_reg=args.lambda_reg, - early_stopping_patience=args.early_stopping, - checkpoint_base=args.checkpoint - ) - - # Save - print(f"\nSaving model to {args.checkpoint}...") - torch.save({ - 'state_dict': {'lt.weight': model.embeddings.weight}, - 'embeddings': model.embeddings.weight, - 'n_nodes': n_nodes, - 'dim': args.dim, - 'max_depth': max_depth, - }, args.checkpoint) - - print("\n" + "="*80) - print("✅ TRAINING COMPLETE") - print("="*80) - print(f"Model saved to: {args.checkpoint}") - print() - print("Next: Run analyze_hierarchy_hyperbolic.py to verify improvements!") - - -if __name__ == "__main__": - main() diff --git a/train_small.py b/train_small.py deleted file mode 100644 index 23185db..0000000 --- a/train_small.py +++ /dev/null @@ -1,470 +0,0 @@ -#!/usr/bin/env python3 -""" -Training script for small dataset with enhanced terminal visualization. - -Features: -- Pre-configured for taxonomy_edges_small dataset -- Real-time metrics display showing improvements -- Compact progress visualization -- Automatic comparison with previous epoch -""" - -import torch -import torch.nn as nn -import torch.optim as optim -import numpy as np -import pandas as pd -import pickle -import argparse -import os -import sys -from collections import defaultdict, deque -from datetime import datetime -from tqdm import tqdm - -# Import the hierarchical model and components -from train_hierarchical import ( - HierarchicalPoincareEmbedding, - HierarchicalDataLoader, - ranking_loss_with_margin, - radial_regularizer -) - - -class MetricsTracker: - """Track and display training metrics with visual improvements.""" - - def __init__(self): - self.history = [] - self.best_loss = float('inf') - self.best_epoch = 0 - - def update(self, epoch, metrics): - """Update metrics for current epoch.""" - self.history.append(metrics) - - if metrics['loss'] < self.best_loss: - self.best_loss = metrics['loss'] - self.best_epoch = epoch - - def get_previous(self, metric_name): - """Get previous epoch's metric value.""" - if len(self.history) < 2: - return None - return self.history[-2].get(metric_name) - - def print_header(self): - """Print column headers.""" - print("\n" + "="*100) - print(f"{'Epoch':>6} | {'Loss':>10} | {'ΔLoss':>10} | {'Improve':>8} | " - f"{'Reg':>8} | {'MaxNorm':>8} | {'Outside':>7} | {'Status':>10}") - print("="*100) - - def print_epoch_summary(self, epoch, metrics, total_epochs): - """Print compact summary of epoch with improvement indicators.""" - prev_loss = self.get_previous('loss') - - # Calculate improvement - if prev_loss is not None: - delta = metrics['loss'] - prev_loss - pct_change = (delta / prev_loss) * 100 if prev_loss != 0 else 0 - - if delta < 0: - status = "✓ BETTER" - delta_str = f"{delta:+.4f}" - pct_str = f"{pct_change:+.2f}%" - improve_color = "\033[92m" # Green - else: - status = "✗ WORSE" - delta_str = f"{delta:+.4f}" - pct_str = f"{pct_change:+.2f}%" - improve_color = "\033[91m" # Red - - reset_color = "\033[0m" - else: - delta_str = "---" - pct_str = "---" - status = "FIRST" - improve_color = "" - reset_color = "" - - # Format output - outside_pct = (metrics['outside_count'] / metrics['total_nodes']) * 100 if metrics['total_nodes'] > 0 else 0 - - print(f"{epoch:6d} | " - f"{metrics['loss']:10.6f} | " - f"{improve_color}{delta_str:>10}{reset_color} | " - f"{improve_color}{pct_str:>8}{reset_color} | " - f"{metrics['reg_loss']:8.6f} | " - f"{metrics['max_norm']:8.4f} | " - f"{outside_pct:6.2f}% | " - f"{improve_color}{status:>10}{reset_color}") - - # Additional info every 5 epochs - if epoch % 5 == 0 or epoch == 1: - print(f" └─ Best: {self.best_loss:.6f} @ epoch {self.best_epoch} | " - f"Norms: [{metrics['min_norm']:.4f}, {metrics['mean_norm']:.4f}, {metrics['max_norm']:.4f}]") - - def print_final_summary(self): - """Print final training summary.""" - print("\n" + "="*100) - print("TRAINING SUMMARY") - print("="*100) - print(f"Total epochs: {len(self.history)}") - print(f"Best loss: {self.best_loss:.6f} (epoch {self.best_epoch})") - - if len(self.history) >= 2: - first_loss = self.history[0]['loss'] - last_loss = self.history[-1]['loss'] - total_improvement = first_loss - last_loss - pct_improvement = (total_improvement / first_loss) * 100 - print(f"Total improvement: {total_improvement:+.6f} ({pct_improvement:+.2f}%)") - - print("="*100 + "\n") - - -def train_with_visualization(model, dataloader, optimizer, n_epochs, - idx_to_depth, max_depth, device, - margin=0.2, lambda_reg=0.1, early_stopping_patience=3, - checkpoint_base=None): - """Train with enhanced terminal visualization.""" - - model.to(device) - model.train() - - print("\n" + "🚀 "*20) - print("HIERARCHICAL TRAINING - SMALL DATASET") - print("🚀 "*20) - print(f"\nConfiguration:") - print(f" Margin: {margin}") - print(f" Regularization: λ={lambda_reg}") - print(f" Early stopping: {'disabled' if early_stopping_patience == 0 else f'{early_stopping_patience} epochs'}") - print(f" Device: {device}") - print(f" Max epochs: {n_epochs}") - - # Precompute regularizer tensors - print("\nPrecomputing regularization targets...") - reg_indices = [] - reg_target_radii = [] - for idx, depth in idx_to_depth.items(): - if idx < model.n_nodes: - reg_indices.append(idx) - target_radius = 0.1 + (depth / max_depth) * 0.85 - reg_target_radii.append(target_radius) - - reg_indices_tensor = torch.LongTensor(reg_indices).to(device) - reg_target_radii_tensor = torch.FloatTensor(reg_target_radii).to(device) - print(f" ✓ Regularizing {len(reg_indices):,} nodes") - - # Metrics tracker - tracker = MetricsTracker() - tracker.print_header() - - # Early stopping - epochs_without_improvement = 0 - checkpoint_queue = deque(maxlen=5) - - for epoch in range(1, n_epochs + 1): - epoch_loss = 0 - epoch_reg_loss = 0 - n_batches = 0 - - # Progress bar for batches - pbar = tqdm(dataloader, - desc=f"Epoch {epoch:3d}/{n_epochs}", - bar_format='{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]', - ncols=100, - leave=False, # Don't leave the bar on screen after completion - dynamic_ncols=True, # Adjust to terminal width - mininterval=0.5) # Update at most every 0.5 seconds - - for ancestors, descendants, negatives, depths in pbar: - ancestors = ancestors.to(device) - descendants = descendants.to(device) - negatives = negatives.to(device) - depths = depths.to(device) - - optimizer.zero_grad() - - # Ranking loss - loss = ranking_loss_with_margin( - model, ancestors, descendants, negatives, - depths, margin=margin, depth_weight=True - ) - - # Radial regularizer - reg_loss = radial_regularizer(model, reg_indices_tensor, - reg_target_radii_tensor, lambda_reg) - - # Total loss - total_loss = loss + reg_loss - total_loss.backward() - - # Gradient clipping - torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) - - optimizer.step() - - # Selective projection - updated_indices = torch.cat([ancestors, descendants, negatives.flatten()]) - updated_indices = torch.unique(updated_indices) - model.project_to_ball(updated_indices) - - # Periodic full projection - if n_batches % 500 == 0: - model.project_to_ball(indices=None) - - epoch_loss += loss.item() - epoch_reg_loss += reg_loss.item() - n_batches += 1 - - # Update progress bar with current batch metrics - pbar.set_postfix({ - 'loss': f'{loss.item():.4f}', - 'reg': f'{reg_loss.item():.4f}' - }) - - # Final projection at epoch end - model.project_to_ball(indices=None) - - # Compute metrics - avg_loss = epoch_loss / n_batches - avg_reg = epoch_reg_loss / n_batches - norms = model.embeddings.weight.norm(dim=1) - outside_count = (norms >= 1.0).sum().item() - - metrics = { - 'loss': avg_loss, - 'reg_loss': avg_reg, - 'min_norm': norms.min().item(), - 'mean_norm': norms.mean().item(), - 'max_norm': norms.max().item(), - 'outside_count': outside_count, - 'total_nodes': model.n_nodes - } - - # Early stopping check - MUST happen BEFORE updating tracker - prev_best_loss = tracker.best_loss - - # Update tracker and display - tracker.update(epoch, metrics) - tracker.print_epoch_summary(epoch, metrics, n_epochs) - - # Save checkpoint - if checkpoint_base: - checkpoint_path = checkpoint_base.replace('.pth', f'_epoch{epoch}.pth') - torch.save({ - 'state_dict': {'lt.weight': model.embeddings.weight}, - 'embeddings': model.embeddings.weight, - 'epoch': epoch, - 'loss': avg_loss, - 'reg_loss': avg_reg, - 'best_loss': tracker.best_loss, - 'epochs_without_improvement': epochs_without_improvement, - }, checkpoint_path) - - # Manage checkpoint queue - if len(checkpoint_queue) >= checkpoint_queue.maxlen: - old_checkpoint = checkpoint_queue[0] - if os.path.exists(old_checkpoint): - os.remove(old_checkpoint) - - checkpoint_queue.append(checkpoint_path) - - # Check if loss improved (compare against PREVIOUS best, not updated best) - if avg_loss < prev_best_loss: - epochs_without_improvement = 0 - - # Save best model - if checkpoint_base: - best_checkpoint = checkpoint_base.replace('.pth', '_best.pth') - torch.save({ - 'state_dict': {'lt.weight': model.embeddings.weight}, - 'embeddings': model.embeddings.weight, - 'epoch': epoch, - 'loss': avg_loss, - 'reg_loss': avg_reg, - }, best_checkpoint) - else: - epochs_without_improvement += 1 - - # Only check early stopping if patience > 0 (0 means disabled) - if early_stopping_patience > 0 and epochs_without_improvement >= early_stopping_patience: - print(f"\n🛑 Early stopping triggered after {epoch} epochs") - print(f" Best loss: {tracker.best_loss:.6f} (epoch {tracker.best_epoch})") - break - - # Final summary - tracker.print_final_summary() - - -def main(): - parser = argparse.ArgumentParser(description='Train on small dataset with enhanced visualization') - parser.add_argument('--data', default='data/taxonomy_edges_small_transitive.pkl', - help='Training data (default: small dataset)') - parser.add_argument('--checkpoint', default='taxonomy_model_small.pth', - help='Output checkpoint path') - parser.add_argument('--mapping', default='data/taxonomy_edges_small.mapping.tsv', - help='Mapping file path aligned with training data') - parser.add_argument('--dim', type=int, default=10, - help='Embedding dimension') - parser.add_argument('--epochs', type=int, default=100, - help='Maximum number of epochs') - parser.add_argument('--early-stopping', type=int, default=5, - help='Early stopping patience') - parser.add_argument('--batch-size', type=int, default=64, - help='Batch size') - parser.add_argument('--n-negatives', type=int, default=50, - help='Number of negative samples') - parser.add_argument('--lr', type=float, default=0.005, - help='Learning rate') - parser.add_argument('--margin', type=float, default=0.2, - help='Ranking loss margin') - parser.add_argument('--lambda-reg', type=float, default=0.1, - help='Radial regularization weight') - parser.add_argument('--gpu', type=int, default=-1, - help='GPU device (-1 for CPU)') - - args = parser.parse_args() - - # Device - if args.gpu >= 0 and torch.cuda.is_available(): - device = torch.device(f'cuda:{args.gpu}') - else: - device = torch.device('cpu') - - # Load training data - print(f"\nLoading training data from {args.data}...") - if not os.path.exists(args.data): - print(f"❌ Error: Training data not found at {args.data}") - print(f"\nPlease run first:") - print(f" python build_transitive_closure.py") - sys.exit(1) - - with open(args.data, 'rb') as f: - training_data = pickle.load(f) - print(f" ✓ Loaded {len(training_data):,} training pairs") - - # Get dataset info - CRITICAL: use mapping file, not training data! - # Training data may not include all nodes (e.g., leaf nodes, isolated nodes) - print(f"Loading mapping to determine true n_nodes from {args.mapping}...") - if not os.path.exists(args.mapping): - print(f"❌ Error: Mapping file not found at {args.mapping}") - sys.exit(1) - - mapping_df = pd.read_csv(args.mapping, sep="\t", dtype=str) - if "taxid" not in mapping_df.columns or "idx" not in mapping_df.columns: - # Re-read assuming headerless file - mapping_df = pd.read_csv( - args.mapping, - sep="\t", - dtype=str, - header=None, - names=["taxid", "idx"], - ) - else: - mapping_df = mapping_df[["taxid", "idx"]] - - mapping_df['idx'] = pd.to_numeric(mapping_df['idx'], errors='coerce') - mapping_df['taxid'] = pd.to_numeric(mapping_df['taxid'], errors='coerce') - mapping_df = mapping_df.dropna(subset=['idx', 'taxid']) - mapping_df['idx'] = mapping_df['idx'].astype(int) - mapping_df['taxid'] = mapping_df['taxid'].astype(int) - n_nodes = int(mapping_df['idx'].max()) + 1 - - max_depth = max(item['descendant_depth'] for item in training_data) - - print(f" Nodes: {n_nodes:,}") - print(f" Max depth: {max_depth}") - - # Build depth mapping from training data - idx_to_depth = {} - for item in training_data: - idx_to_depth[item['descendant_idx']] = item['descendant_depth'] - if item['ancestor_idx'] not in idx_to_depth: - idx_to_depth[item['ancestor_idx']] = item['ancestor_depth'] - - # For nodes NOT in training data, load their depths from full taxonomy - nodes_in_training = set(idx_to_depth.keys()) - missing_nodes = n_nodes - len(nodes_in_training) - - if missing_nodes > 0: - print(f"\n⚠️ {missing_nodes:,} nodes not in training data - loading their depths from taxonomy...") - - # Build TaxID -> depth mapping from training data - taxid_to_depth = {} - for item in training_data: - taxid_to_depth[item['ancestor_taxid']] = item['ancestor_depth'] - taxid_to_depth[item['descendant_taxid']] = item['descendant_depth'] - - # Map missing indices to depths via TaxID lookup - for idx in range(n_nodes): - if idx not in idx_to_depth: - # Find TaxID for this index - taxid = mapping_df[mapping_df['idx'] == idx]['taxid'].values[0] - if taxid in taxid_to_depth: - idx_to_depth[idx] = taxid_to_depth[taxid] - else: - # Node not in taxonomy at all - likely a leaf, assign max depth - idx_to_depth[idx] = max_depth - - print(f" ✓ Assigned depths to {missing_nodes:,} additional nodes") - print(f" ✓ Total nodes with depth info: {len(idx_to_depth):,} / {n_nodes:,}") - - # Create model - print("\nInitializing model with depth-aware embeddings...") - model = HierarchicalPoincareEmbedding( - n_nodes=n_nodes, - dim=args.dim, - max_depth=max_depth, - init_depth_data=idx_to_depth - ) - - # Create dataloader - print("Creating dataloader with hard negative sampling...") - dataloader = HierarchicalDataLoader( - training_data=training_data, - n_nodes=n_nodes, - batch_size=args.batch_size, - n_negatives=args.n_negatives, - depth_stratify=True - ) - - # Optimizer - optimizer = optim.Adam(model.parameters(), lr=args.lr) - - # Train with visualization - train_with_visualization( - model=model, - dataloader=dataloader, - optimizer=optimizer, - n_epochs=args.epochs, - idx_to_depth=idx_to_depth, - max_depth=max_depth, - device=device, - margin=args.margin, - lambda_reg=args.lambda_reg, - early_stopping_patience=args.early_stopping, - checkpoint_base=args.checkpoint - ) - - # Save final model - print(f"Saving final model to {args.checkpoint}...") - torch.save({ - 'state_dict': {'lt.weight': model.embeddings.weight}, - 'embeddings': model.embeddings.weight, - 'n_nodes': n_nodes, - 'dim': args.dim, - 'max_depth': max_depth, - }, args.checkpoint) - - print("\n✅ Training complete!") - print(f" Model saved: {args.checkpoint}") - print(f" Best model: {args.checkpoint.replace('.pth', '_best.pth')}") - print("\nNext steps:") - print(" python analyze_hierarchy_hyperbolic.py") - print(f" python scripts/visualize_embeddings.py {args.checkpoint.replace('.pth', '_best.pth')}") - - -if __name__ == "__main__": - main() diff --git a/uv.lock b/uv.lock index 913da9f..aa32873 100644 --- a/uv.lock +++ b/uv.lock @@ -1217,7 +1217,7 @@ wheels = [ [[package]] name = "taxembed" -version = "0.2.0" +version = "1.0.0" source = { editable = "." } dependencies = [ { name = "matplotlib" }, @@ -1238,14 +1238,6 @@ dev = [ { name = "ruff" }, ] -[package.dev-dependencies] -dev = [ - { name = "mypy" }, - { name = "pytest" }, - { name = "pytest-cov" }, - { name = "ruff" }, -] - [package.metadata] requires-dist = [ { name = "matplotlib", specifier = ">=3.5.0" }, @@ -1263,14 +1255,6 @@ requires-dist = [ ] provides-extras = ["dev"] -[package.metadata.requires-dev] -dev = [ - { name = "mypy", specifier = ">=1.0.0" }, - { name = "pytest", specifier = ">=8.0.0" }, - { name = "pytest-cov", specifier = ">=4.1.0" }, - { name = "ruff", specifier = ">=0.6.0" }, -] - [[package]] name = "taxopy" version = "0.14.0"