Epic: htfa-parameter-support

[jeremymanning]

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