# VSAX: Vector Symbolic Algebra for JAX

[![PyPI version](https://img.shields.io/pypi/v/vsax.svg)](https://pypi.org/project/vsax/)
[![Python Version](https://img.shields.io/pypi/pyversions/vsax.svg)](https://pypi.org/project/vsax/)
[![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)
[![Documentation](https://img.shields.io/badge/docs-github%20pages-blue)](https://vasanthsarathy.github.io/vsax/)
[![Code style: ruff](https://img.shields.io/badge/code%20style-ruff-000000.svg)](https://github.com/astral-sh/ruff)

VSAX is a GPU-accelerated, JAX-native Python library for Vector Symbolic Architectures (VSAs). It provides composable symbolic representations using hypervectors, algebraic operations for binding and bundling, and encoding strategies for symbolic and structured data.

## 📚 New to VSA? Start Learning!

**[Complete VSA Learning Path →](https://vasanthsarathy.github.io/vsax/course/)**

New comprehensive course covering VSA from foundations to research frontiers:

- **20 Interactive Lessons** - From "Why High Dimensions?" to research frontiers
- **12 Hands-On Exercises** - ~5,000 lines of practical code
- **5 Progressive Modules** - Foundations → Core Operations → Applications → Advanced → Research
- **~113,000 Words** of content - Self-paced learning for ML practitioners and researchers

**Quick Links:**
- [Module 1: Foundations](https://vasanthsarathy.github.io/vsax/course/01_foundations/) - Why VSA works
- [Module 2: Core Operations](https://vasanthsarathy.github.io/vsax/course/02_operations/) - FHRR, MAP, Binary models
- [Module 3: Encoders & Applications](https://vasanthsarathy.github.io/vsax/course/03_encoders/) - Real-world usage
- [Module 4: Advanced Techniques](https://vasanthsarathy.github.io/vsax/course/04_advanced/) - State-of-the-art methods
- [Module 5: Research & Extensions](https://vasanthsarathy.github.io/vsax/course/05_research/) - Contributing to VSA research

## Features

- 🚀 **Three VSA Models**: FHRR, MAP, and Binary implementations ✅
- 🏭 **Factory Functions**: One-line model creation with sensible defaults ✅
- 💾 **VSAMemory**: Dictionary-style symbol management ✅
- 📊 **5 Core Encoders**: Scalar, Sequence, Set, Dict, and Graph encoders ✅
- 🎨 **Custom Encoders**: Easy-to-extend AbstractEncoder base class ✅
- 🔍 **Similarity Metrics**: Cosine, dot, and Hamming similarity ✅
- ⚡ **Batch Operations**: GPU-accelerated vmap operations for parallel processing ✅
- 💾 **I/O & Persistence**: Save/load basis vectors to JSON ✅
- 🎮 **GPU Utilities**: Device management, benchmarking, CPU/GPU comparison ✅
- 🔧 **Clifford Operators**: Exact, compositional, invertible transformations for reasoning ✅ **NEW in v1.1.0**
- 📍 **Fractional Power Encoding (FPE)**: Continuous value encoding using v^r for spatial/function representations ✅ **NEW in v1.2.0**
- 🗺️ **Spatial Semantic Pointers (SSP)**: Continuous spatial encoding with what/where queries ✅ **NEW in v1.2.0**
- 🎯 **Vector Function Architecture (VFA)**: Function approximation in RKHS for density estimation, regression, image processing ✅ **NEW in v1.2.0**
- 🚀 **GPU-Accelerated**: Built on JAX for automatic GPU acceleration (5-30x speedup)
- 🧩 **Modular Architecture**: Clean separation between representations and operations
- 🧬 **Complete Representations**: Complex, Real, and Binary hypervectors ✅
- ⚙️ **Full Operation Sets**: FFT-based FHRR, MAP, and XOR/majority Binary ops ✅
- 🎲 **Random Sampling**: Sampling utilities for all representation types ✅
- 📚 **Comprehensive Documentation**: Full API docs and examples ✅
- 📓 **Interactive Tutorials**: Jupyter notebooks with real datasets (MNIST, knowledge graphs) ✅ **NEW in v0.7.1**
- ✅ **95% Test Coverage**: 450 tests ensuring reliability

## Installation

### From PyPI (Recommended)

```bash
pip install vsax
```

Or with uv:

```bash
uv pip install vsax
```

### From Source

#### Using uv (Recommended)

[uv](https://github.com/astral-sh/uv) is a fast Python package installer and resolver. Install it first:

```bash
# Install uv (Unix/macOS)
curl -LsSf https://astral.sh/uv/install.sh | sh

# Install uv (Windows)
powershell -c "irm https://astral.sh/uv/install.ps1 | iex"
```

Then install VSAX:

```bash
git clone https://github.com/vasanthsarathy/vsax.git
cd vsax

# Create virtual environment and install package
uv venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate
uv pip install -e .
```

#### Using pip

```bash
git clone https://github.com/vasanthsarathy/vsax.git
cd vsax
pip install -e .
```

### Development Installation

#### Using uv (Recommended)

```bash
uv venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate
uv pip install -e ".[dev,docs]"
```

#### Using pip

```bash
pip install -e ".[dev,docs]"
```

## Quick Start

**New in v0.6.0:** Save and load basis vectors!

### Simple Example

```python
from vsax import create_fhrr_model, VSAMemory, DictEncoder, save_basis, load_basis
from vsax.similarity import cosine_similarity
from vsax.utils import vmap_bind

# Create model with factory function (one line!)
model = create_fhrr_model(dim=512)

# Create memory for symbol management
memory = VSAMemory(model)

# Add symbols - automatically samples and stores hypervectors
memory.add_many(["subject", "action", "dog", "run", "cat", "jump"])

# Dictionary-style access
dog = memory["dog"]

# Encode structured data with DictEncoder
encoder = DictEncoder(model, memory)
sentence = encoder.encode({"subject": "dog", "action": "run"})

# Bind concepts (circular convolution)
dog_is_animal = model.opset.bind(dog.vec, memory["animal"].vec)

# Unbind to recover original concept (NEW: explicit unbind method)
recovered = model.opset.unbind(dog_is_animal, memory["animal"].vec)
# With FHRR: >99% similarity to original dog vector

# Bundle concepts (sum and normalize)
pets = model.opset.bundle(memory["dog"].vec, memory["cat"].vec)

# NEW: Similarity search
similarity = cosine_similarity(memory["dog"], memory["cat"])
print(f"Dog-Cat similarity: {similarity:.3f}")

# NEW: Batch operations (GPU-accelerated)
import jax.numpy as jnp
nouns = jnp.stack([memory["dog"].vec, memory["cat"].vec])
verbs = jnp.stack([memory["run"].vec, memory["jump"].vec])
actions = vmap_bind(model.opset, nouns, verbs)  # Parallel binding!

# NEW: Save and load basis vectors
save_basis(memory, "my_basis.json")  # Persist to JSON
memory_new = VSAMemory(model)
load_basis(memory_new, "my_basis.json")  # Load from JSON
```

### All Three Models

VSAX supports three VSA models, all with the same simple API:

```python
from vsax import create_fhrr_model, create_map_model, create_binary_model, VSAMemory

# FHRR: Complex hypervectors, >99% unbinding accuracy (with proper sampling)
fhrr = create_fhrr_model(dim=512)

# MAP: Real hypervectors, approximate unbinding
map_model = create_map_model(dim=512)

# Binary: Discrete hypervectors, exact unbinding (XOR self-inverse)
binary = create_binary_model(dim=10000, bipolar=True)

# Same interface for all models!
for model in [fhrr, map_model, binary]:
    memory = VSAMemory(model)
    memory.add("concept")
    vec = memory["concept"]
```

### NEW: Clifford Operators (v1.1.0)

**Exact, compositional, invertible transformations for reasoning:**

```python
from vsax import create_fhrr_model, VSAMemory
from vsax.operators import create_left_of, create_agent, create_patient
from vsax.similarity import cosine_similarity

model = create_fhrr_model(dim=512)
memory = VSAMemory(model)
memory.add_many(["cup", "plate", "dog", "cat", "chase"])

# Use pre-defined spatial operators
LEFT_OF = create_left_of(512)

# Encode: cup LEFT_OF plate
scene = model.opset.bundle(
    memory["cup"].vec,
    LEFT_OF.apply(memory["plate"]).vec
)

# Query with inverse operator
RIGHT_OF = LEFT_OF.inverse()
answer = LEFT_OF.inverse().apply(model.rep_cls(scene))
# Returns plate with high similarity!

# Use pre-defined semantic operators
AGENT = create_agent(512)
PATIENT = create_patient(512)

# Encode: "dog chases cat"
sentence = model.opset.bundle(
    AGENT.apply(memory["dog"]).vec,
    memory["chase"].vec,
    PATIENT.apply(memory["cat"]).vec
)

# Query: Who is the AGENT?
who = AGENT.inverse().apply(model.rep_cls(sentence))
similarity = cosine_similarity(who.vec, memory["dog"].vec)
print(f"AGENT is 'dog': {similarity:.3f}")  # High similarity!
```

**Pre-defined operators (NEW in Phase 2):**
- **Spatial**: `create_left_of`, `create_right_of`, `create_above`, `create_below`, `create_in_front_of`, `create_behind`, `create_near`, `create_far`
- **Semantic**: `create_agent`, `create_patient`, `create_theme`, `create_experiencer`, `create_instrument`, `create_location`, `create_goal`, `create_source`

**Key features:**
- ✅ **Exact inversion**: `op.inverse().apply(op.apply(v))` recovers original (similarity > 0.999)
- ✅ **Explicit unbinding**: Use `ops.unbind(bound, b)` for clearer unbinding operations
- ✅ **Compositional**: Combine operators algebraically with `compose()`
- ✅ **Typed**: Semantic metadata (SPATIAL, SEMANTIC, TEMPORAL, etc.)
- ✅ **Reproducible**: Same dimension always produces same operator
- ✅ **FHRR-compatible**: Phase-based transformations for complex hypervectors

### NEW: Spatial Semantic Pointers (v1.2.0)

**Continuous spatial encoding for 2D navigation and spatial reasoning:**

```python
from vsax import create_fhrr_model, VSAMemory
from vsax.spatial import SpatialSemanticPointers, SSPConfig

model = create_fhrr_model(dim=1024)
memory = VSAMemory(model)

# Create 2D spatial encoder
config = SSPConfig(dim=1024, num_axes=2, axis_names=["X", "Y"])
ssp = SpatialSemanticPointers(model, memory, config)

# Add objects to memory
memory.add_many(["table", "chair", "plant"])

# Bind objects to locations in a 10x10 room
table_at_loc = ssp.bind_object_location("table", [3.0, 7.0])
chair_at_loc = ssp.bind_object_location("chair", [5.0, 7.5])
plant_at_loc = ssp.bind_object_location("plant", [8.0, 3.0])

# Bundle into scene
from vsax.representations import ComplexHypervector
scene_vec = model.opset.bundle(
    table_at_loc.vec,
    model.opset.bundle(chair_at_loc.vec, plant_at_loc.vec)
)
scene = ComplexHypervector(scene_vec)

# Query: What's at location (3.0, 7.0)?
result = ssp.query_location(scene, [3.0, 7.0])
# Returns hypervector similar to "table"!

# Query: Where is the plant?
location_hv = ssp.query_object(scene, "plant")
decoded_location = ssp.decode_location(
    location_hv,
    search_range=[(0, 10), (0, 10)],
    resolution=20
)
# Returns approximately [8.0, 3.0]!
```

**Key features:**
- ✅ **Continuous encoding**: Positions encoded as X^x ⊗ Y^y using fractional powers
- ✅ **What/Where queries**: Both "what is at (x,y)" and "where is X" supported
- ✅ **Scene shifting**: Transform entire scenes spatially
- ✅ **Similarity maps**: Visualize spatial distributions

### NEW: Vector Function Architecture (v1.2.0)

**Function approximation in RKHS for density estimation, regression, and image processing:**

```python
from vsax import create_fhrr_model, VSAMemory
from vsax.vfa import DensityEstimator, NonlinearRegressor, ImageProcessor
import jax.numpy as jnp

model = create_fhrr_model(dim=2048)

# 1. Density Estimation
estimator = DensityEstimator(model, VSAMemory(model), bandwidth=0.5)
samples = jax.random.normal(jax.random.PRNGKey(42), (100,))
estimator.fit(samples)
density = estimator.evaluate(jnp.array([0.0, 1.0, 2.0]))

# 2. Nonlinear Regression
regressor = NonlinearRegressor(model, VSAMemory(model))
x_train = jnp.linspace(0, 2 * jnp.pi, 50)
y_train = jnp.sin(x_train)
regressor.fit(x_train, y_train)
y_pred = regressor.predict(jnp.array([0.5, 1.0, 1.5]))
r2_score = regressor.score(x_train, y_train)

# 3. Image Processing
processor = ImageProcessor(model, VSAMemory(model))
image = jnp.zeros((16, 16))
image = image.at[4:12, 4:12].set(1.0)  # White square
image_hv = processor.encode(image)
reconstructed = processor.decode(image_hv, shape=(16, 16))

# Blend two images
image2 = jnp.zeros((16, 16))
image2_hv = processor.encode(image2)
blended = processor.blend(image_hv, image2_hv, alpha=0.5)
```

**Key features:**
- ✅ **Kernel density estimation**: Estimate probability densities from samples
- ✅ **Nonlinear regression**: Fit complex functions with scikit-learn-like API
- ✅ **Image encoding**: Compact hypervector representation of images
- ✅ **Function arithmetic**: Add, scale, shift, and convolve functions
- ✅ **RKHS approximation**: Reproducing Kernel Hilbert Space function representation

### Advanced: Manual Model Creation

You can still create models manually if you need custom configuration:

```python
from vsax import VSAModel, ComplexHypervector, FHRROperations, sample_complex_random

model = VSAModel(
    dim=512,
    rep_cls=ComplexHypervector,
    opset=FHRROperations(),
    sampler=sample_complex_random
)
```

See [docs/design-spec.md](docs/design-spec.md) for complete technical specification.

## Development Status

Currently in **Iteration 6: I/O & Persistence** ✅

### Completed

**Iteration 1** (v0.1.0): Foundation & Infrastructure ✅
- ✅ Core abstract classes (AbstractHypervector, AbstractOpSet)
- ✅ VSAModel dataclass
- ✅ Package structure
- ✅ Testing infrastructure (pytest, coverage)
- ✅ CI/CD pipeline (GitHub Actions)
- ✅ Documentation site (MkDocs)
- ✅ Development tooling (ruff, mypy)

**Iteration 2** (v0.2.0): All 3 Representations + All 3 OpSets ✅
- ✅ ComplexHypervector, RealHypervector, BinaryHypervector
- ✅ FHRROperations, MAPOperations, BinaryOperations
- ✅ Sampling utilities (sample_random, sample_complex_random, sample_binary_random)
- ✅ 175 comprehensive tests with 96% coverage
- ✅ Full integration tests for all model combinations

**Iteration 3** (v0.3.0): VSAMemory + Factory Functions ✅
- ✅ VSAMemory class - dictionary-style symbol management
- ✅ Factory functions (create_fhrr_model, create_map_model, create_binary_model)
- ✅ Utility functions (coerce_to_array, validation helpers)
- ✅ 230 tests with 89% coverage
- ✅ Comprehensive documentation guides

**Iteration 4** (v0.4.0): Encoders + Examples ✅ **FIRST USABLE RELEASE!**
- ✅ ScalarEncoder - Numeric values with power encoding
- ✅ SequenceEncoder - Ordered sequences (lists, tuples)
- ✅ SetEncoder - Unordered collections (sets)
- ✅ DictEncoder - Key-value pairs (dictionaries)
- ✅ GraphEncoder - Graph structures (edge lists)
- ✅ AbstractEncoder - Base class for custom encoders
- ✅ Complete integration examples for all 3 models
- ✅ Custom encoder examples (DateEncoder, ColorEncoder)
- ✅ 280+ tests with 92%+ coverage

**Iteration 5** (v0.5.0): Similarity Metrics & Utilities ✅
- ✅ Cosine, dot, and Hamming similarity functions
- ✅ Batch operations with JAX vmap (vmap_bind, vmap_bundle, vmap_similarity)
- ✅ Visualization utilities (pretty_repr, format_similarity_results)
- ✅ GPU-accelerated similarity search
- ✅ Comprehensive examples (similarity_search.py, batch_operations.py)
- ✅ 319 tests with 95%+ coverage

**Iteration 6** (v0.6.0): I/O & Persistence ✅
- ✅ save_basis() and load_basis() functions
- ✅ JSON serialization for all 3 models
- ✅ Round-trip vector preservation
- ✅ Dimension and type validation
- ✅ Comprehensive tests (339 tests, 96% coverage)
- ✅ Complete examples and documentation

### Coming Next

**Iteration 7** (v1.0.0): Full Documentation & Production Release
- Complete API documentation
- Tutorial notebooks
- Production-ready v1.0.0 release

See [todo.md](todo.md) for the complete development roadmap.

## Documentation

- [Full Documentation](https://vasanthsarathy.github.io/vsax/)
- [Getting Started](docs/getting-started.md)
- [Modeling Guide](docs/modeling-guide.md) - **7-step workflow** for building VSA applications
- [Tutorials](docs/tutorials/index.md)
  - [MNIST Classification](docs/tutorials/01_mnist_classification.md) - Image classification with VSA
  - [Knowledge Graph Reasoning](docs/tutorials/02_knowledge_graph.md) - Multi-hop reasoning with relational facts
  - [Kanerva's "Dollar of Mexico"](docs/tutorials/03_kanerva_analogies.md) - Analogical reasoning and mapping
  - [Word Analogies & Random Indexing](docs/tutorials/04_word_analogies.md) - Word embeddings and semantic similarity
  - [Understanding VSA Models](docs/tutorials/05_model_comparison.md) - Compare FHRR, MAP, and Binary models
  - [VSA for Edge Computing](docs/tutorials/06_edge_computing.md) - Lightweight alternative to neural networks
  - [Hierarchical Structures](docs/tutorials/07_hierarchical_structures.md) - Trees and nested composition
  - [Multi-Modal Concept Grounding](docs/tutorials/08_multimodal_grounding.md) - Fuse vision, language, and arithmetic
  - [Neural-Symbolic Fusion (HD-Glue)](docs/tutorials/09_neural_symbolic_fusion.md) - Glue neural networks with VSA ✨ **NEW**
- [VSAX vs Other Libraries](docs/comparison.md) - Why VSAX? Feature comparison with torchhd, hdlib, PyBHV
- [Design Specification](docs/design-spec.md)
- [API Reference](https://vasanthsarathy.github.io/vsax/api/)
- [Contributing](CONTRIBUTING.md)

## Contributing

We welcome contributions! Please see [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines.

## Development

Run tests:
```bash
pytest
```

Run tests with coverage:
```bash
pytest --cov=vsax --cov-report=term-missing
```

Type checking:
```bash
mypy vsax
```

Linting:
```bash
ruff check vsax tests
```

Build documentation:
```bash
mkdocs serve
```

## License

VSAX is released under the MIT License. See [LICENSE](LICENSE) for details.

## Citation

If you use VSAX in your research, please cite:

```bibtex
@software{vsax2025,
  title = {VSAX: Vector Symbolic Algebra for JAX},
  author = {Sarathy, Vasanth},
  year = {2025},
  url = {https://github.com/vasanthsarathy/vsax},
  version = {1.3.0}
}
```

## Acknowledgments

VSAX is built on [JAX](https://github.com/google/jax) and inspired by the VSA/HDC research community.