ContextLab / ContextLab/htfa

Epic: htfa-parameter-support

Open
#155 0 comments 0 reactions 0 assignees View on GitHub
epic epic:htfa-parameter-support feature
Dominant language
Python
Stars
1
Forks
0
PR merge metrics
No merged PRs in 30d

Description

# Epic: htfa-parameter-support

## Overview
Extend the HTFA implementation with scikit-learn compliant parameters (n_levels, backend, random_state, max_iter) while maintaining backward compatibility. Focus on a pragmatic NumPy-first implementation with extensible backend architecture for future GPU acceleration.

## Architecture Decisions

### Parameter Strategy
- **Incremental Enhancement**: Add parameters to existing `__init__` without modifying current behavior
- **Smart Defaults**: `backend=None` triggers automatic selection based on available resources
- **Graceful Degradation**: Automatic fallback to NumPy when specialized backends unavailable

### Backend Architecture
- **Strategy Pattern**: Abstract `Backend` interface with swappable implementations
- **Lazy Loading**: Only import JAX/PyTorch when explicitly requested
- **Resource Detection**: Automatic check for GPU/Metal/CUDA availability when backend=None

### Random State Management
- **Centralized Control**: Single `RandomState` object propagated to all stochastic operations
- **Scikit-learn Compliance**: Support int seed, RandomState instance, or None
- **Existing Integration**: Leverage existing `check_random_state` from sklearn.utils

## Technical Approach

### Core Components

#### 1. Parameter Extension (`htfa/core/htfa.py`)
- Extend `HTFA.__init__` with new parameters
- Add validation using existing `htfa/validation.py` framework
- Ensure defaults preserve current behavior

#### 2. Backend Interface (`htfa/backends/base.py`)
- Abstract base class defining numerical operations
- Methods: matrix ops, decompositions, random sampling
- Automatic backend selection logic

#### 3. NumPy Backend (`htfa/backends/numpy_backend.py`)
- Default implementation using existing NumPy code
- Minimal refactoring - wrap existing operations
- Baseline for performance comparisons

#### 4. Random State Propagation
- Update K-means initialization in `htfa/core/tfa.py`
- Modify factor initialization routines
- Pass through to scipy/sklearn functions

#### 5. Convergence Control
- Add iteration counter to optimization loops
- Implement early stopping based on factor change magnitude
- Expose convergence metrics in results

### Integration Points
- **Existing TFA Class**: Minimal changes, mainly random_state threading
- **Validation Framework**: Extend with new parameter validators
- **Test Suite**: Re-enable disabled tests progressively

## Implementation Strategy

### Phase 1: Core Parameters (Priority)
1. Add parameters with validation
2. Thread random_state through existing code
3. Implement max_iter in optimization loops
4. Verify backward compatibility

### Phase 2: Backend Foundation
1. Create minimal backend interface
2. Wrap existing NumPy code as default backend
3. Add backend selection logic
4. Document extension points for JAX/PyTorch

### Phase 3: Hierarchical Levels
1. Refactor fixed 2-level assumption
2. Support n_levels=1 (standard FA mode)
3. Validate level-data compatibility

### Testing Approach
- **Progressive Enablement**: Re-enable tests one by one as features are implemented. Ensure tests pass before marking features as complete.
- **Parametrized Testing**: Cover parameter combinations systematically
- **Regression Testing**: Ensure identical results with default parameters

## Stats

Total tasks: 8
Parallel tasks: 2 (can be worked on simultaneously)
Sequential tasks: 6 (have dependencies)
Estimated total effort: 46-60 hours

## Dependencies

### Internal
- `htfa/core/htfa.py` - Main implementation file
- `htfa/core/tfa.py` - Base class requiring random_state support
- `htfa/validation.py` - Extend with parameter validators
- `tests/test_htfa_core.py.disabled` - Tests to progressively enable

### External
- NumPy >= 2.0 (already required)
- scikit-learn >= 1.0 (already required)
- JAX/PyTorch (optional, defer to follow-up issue)

### Prerequisites
- No blocking dependencies - can start immediately
- Defer JAX/PyTorch backends to follow-up PR

## Success Criteria (Technical)

### Functionality
- All 4 parameters functional with validation
- Disabled tests passing
- Zero breaking changes to existing API

### Performance
- No regression with default parameters
- Backend selection works correctly
- Random state ensures reproducibility

### Quality Gates
- 100% backward compatibility
- Test coverage maintained above 90%
- All docstrings updated
- Type hints complete

## Estimated Effort

### Overall Timeline
- **Total Duration**: 5-7 days
- **Critical Path**: Parameter addition → Random state → Test enablement

### Effort Distribution
- Parameter implementation: 2 days
- Backend architecture: 1 day
- Hierarchical levels: 1 day
- Testing & documentation: 2 days
- Buffer for issues: 1 day

### Resource Requirements
- Single developer
- Review from maintainer
- No external dependencies

## Risk Mitigation

### Simplified Approach
- Start with NumPy-only implementation
- Defer complex backends to follow-up
- Use existing validation framework
- Minimal refactoring of core logic

### Incremental Delivery
- Each task independently testable
- Progressive test enablement
- Early validation of approach

## Notes

### Simplifications from PRD
- Defer JAX/PyTorch to follow-up issue (document requirements)
- Focus on NumPy backend initially
- Leverage existing code maximally
- Aim for minimal, clean implementation

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.