REVIEW_AND_IMPROVEMENT_PLAN.md

Math Module: Review and Improvement Plan

Warning

Stale Report: This review was conducted on 2025-01-06. Several issues noted below have since been resolved (e.g., fisher_exact_test() now has a real implementation, README.md module references are corrected). The findings may no longer reflect the current state.

Date: 2025-01-06 Scope: Complete review of src/metainformant/math/ module Status: Historical — many items resolved

Executive Summary

The math module is well-structured with comprehensive coverage of theoretical biology topics. It demonstrates strong mathematical rigor and good code organization. However, there are opportunities for improvement in documentation consistency, error handling, test coverage gaps, and API completeness.

Overall Assessment: (4/5) - Strong foundation with clear improvement pathways

---

1. Code Quality Assessment

1.1 Strengths

Modular Design: Clear separation of concerns across submodules Mathematical Rigor: Functions are well-grounded in literature with references Type Hints: Comprehensive type annotations throughout Docstrings: Most functions have detailed docstrings with examples Test Coverage: Extensive test suite (22+ test files) Edge Case Handling: Most functions handle invalid inputs gracefully

1.2 Issues Identified

Critical Issues

1. Placeholder Implementation (__init__.py:305-348)

   • fisher_exact_test() returns placeholder p-value (1.0)
   • Should either implement properly or remove/warn users

2. Missing Exports (__init__.py)

   • Functions defined in __init__.py not in __all__:
     • correlation_coefficient
     • linear_regression
     • fisher_exact_test
     • shannon_entropy
     • jensen_shannon_divergence
   • wattersons_theta defined in coalescent.py but not exported

3. Inconsistent Error Handling

   • Some functions raise ValueError (e.g., effective_population_size_from_heterozygosity)
   • Others return 0.0 or default values
   • No consistent error handling strategy

4. Optional Dependencies

   • scipy used conditionally but not declared as optional dependency
   • numpy used in popgen_stats.py but dependency not clearly documented

Moderate Issues

5. Deprecation Warnings (coalescent.py:632-704)

   • expected_segregating_sites() has deprecated parameter order
   • Should be removed or clearly documented

6. Incomplete Functionality

   • fay_wu_h() doesn't fully implement derived allele counting (lines 347-399)
   • hardy_weinberg_test() has scipy fallback but could be more robust

7. Documentation Inconsistencies

   • Some functions have detailed examples, others don't
   • README.md references non-existent modules (kin_selection.py, multilevel_selection.py)
   • AGENTS.md could be more detailed

8. Code Duplication

   • watterson_theta() exists in both popgen.py and coalescent.py
   • wattersons_theta() also in coalescent.py (alias)
   • Could be consolidated

---

2. Documentation Review

2.1 Module Documentation (__init__.py)

Current State: Good module-level docstring
Issues:

• Missing list of key functions by category
• No usage examples in module docstring
• Doesn't mention optional dependencies

Recommendation: Expand module docstring with:

• Brief overview of each submodule
• Key use cases
• Dependency information

2.2 README.md

Current State: Comprehensive but has inaccuracies
Issues:

• References kin_selection.py and multilevel_selection.py (should be selection.py)
• Missing documentation for __init__.py utility functions
• No troubleshooting section
• Missing performance benchmarks/guidelines

Recommendation:

• Fix module references
• Add section on utility functions in __init__.py
• Add troubleshooting/common issues section
• Document performance characteristics

2.3 AGENTS.md

Current State: Basic, follows template
Issues:

• Generic descriptions
• Doesn't detail specific AI contributions per function
• Missing future enhancement plans

Recommendation: Enhance with:

• Specific examples of AI-assisted implementations
• Function-level contribution notes
• Future AI integration plans

2.4 Function Docstrings

Current State: Generally excellent
Issues:

• Some functions missing examples (e.g., CoalescentSummary methods)
• Inconsistent format for examples
• Some references missing DOIs

Recommendation: Standardize docstring format:

• Always include examples
• Use consistent example format
• Add DOI links where available

---

3. Test Coverage Analysis

3.1 Coverage by Submodule

Submodule	Test Files	Coverage	Status
coalescent.py	3 files	Excellent	[DONE]
price.py	1 file	Good	[DONE]
popgen.py	2 files	Good	[DONE]
ld.py	2 files	Good	[DONE]
epidemiology.py	3 files	Good	[DONE]
dynamics.py	1 file	Good	[DONE]
ddm.py	1 file	Good	[DONE]
selection.py	1 file	Good	[DONE]
quantgen.py	1 file	Good	[DONE]
egt.py	Part of combined	Good	[DONE]
fst.py	Part of combined	Good	[DONE]
effective_size.py	1 file	Good	[DONE]
demography.py	1 file	Good	[DONE]
popgen_stats.py	1 file	Good	[DONE]
__init__.py utilities	1 file	Partial
selection_experiments/	1 file	Good	[DONE]

3.2 Missing Test Coverage

Critical Gaps:

1. __init__.py utility functions:

   • correlation_coefficient() - No dedicated tests
   • linear_regression() - No dedicated tests
   • fisher_exact_test() - Needs tests (placeholder)
   • shannon_entropy() - No dedicated tests
   • jensen_shannon_divergence() - No dedicated tests

2. Edge Cases:

   • Very large sample sizes
   • Extreme parameter values
   • Numerical precision limits

3. Integration Tests:

   • Cross-module functionality
   • Workflow examples

Recommendation: Add tests for:

• All __init__.py utility functions
• Edge cases with extreme values
• Integration scenarios

---

4. API Design Review

4.1 Function Signatures

Strengths:

• Consistent use of type hints
• Clear parameter names
• Good use of keyword-only arguments where appropriate

Issues:

• Some functions have too many parameters (e.g., lotka_volterra_step())
• Inconsistent ordering of similar parameters across functions
• Missing validation for some parameter ranges

Recommendation:

• Consider dataclasses for functions with 5+ parameters
• Standardize parameter ordering (e.g., always sample_size before effective_population_size)
• Add range validation with clear error messages

4.2 Return Types

Strengths:

• Clear return type annotations
• Consistent tuple ordering

Issues:

• Some functions return 0.0 on error, others raise exceptions
• Inconsistent handling of invalid inputs

Recommendation: Establish error handling policy:

• Use exceptions for programming errors (invalid types)
• Use sentinel values (0.0, None) for domain errors (invalid ranges)
• Document error behavior clearly

4.3 Export Completeness

Missing from __all__:

• All utility functions in __init__.py
• wattersons_theta (alias)
• CoalescentSummary class
• site_frequency_spectrum_counts
• expected_sfs_counts
• expected_coalescent_waiting_times

Recommendation: Review all public functions and add to __all__

---

5. Performance Considerations

5.1 Current Performance

Strengths:

• Vectorized operations where appropriate (NumPy)
• Efficient algorithms (e.g., harmonic means)

Issues:

• Some functions use list comprehensions where NumPy would be faster
• No caching for expensive computations (e.g., tajima_constants)
• Some functions could benefit from memoization

Recommendation:

• Profile critical functions
• Add caching for expensive computations (e.g., tajima_constants)
• Consider NumPy vectorization for loops over large datasets

5.2 Memory Usage

Current State: Generally efficient
Recommendation: Document memory requirements for large-scale operations

---

6. Mathematical Correctness

6.1 Formula Verification

Status: Functions appear mathematically correct Recommendation: Add numerical validation tests against known analytical solutions

6.2 Numerical Stability

Issues:

• Some functions may have numerical precision issues with extreme values
• No explicit handling of floating-point errors

Recommendation:

• Add numerical stability tests
• Document precision limitations
• Consider using math.isclose() for comparisons

---

7. Selection Experiments Submodule

7.1 Current State

Strengths:

• Well-documented with comprehensive README
• Clear separation of concerns (model, plotting, CLI)
• Good examples

Issues:

• No integration with main module exports
• CLI could have better error handling
• Missing tests for plotting functions

Recommendation:

• Add selection_experiments to main module exports (optional)
• Improve CLI error messages
• Add tests for plotting functions

---

8. Improvement Plan

Phase 1: Critical Fixes (Priority: High)

1. Fix Placeholder Implementation (1-2 hours)

   • Implement proper fisher_exact_test() or mark as TODO
   • Add warning/docstring about placeholder status

2. Complete Exports (30 minutes)

   • Add all __init__.py functions to __all__
   • Add missing coalescent functions to exports
   • Verify all public APIs are exported

3. Error Handling Consistency (2-3 hours)

   • Establish error handling policy
   • Standardize error responses across module
   • Document error behavior

4. Fix Documentation Inaccuracies (1 hour)

   • Update README.md module references
   • Fix broken cross-references
   • Add missing function documentation

Phase 2: Enhancements (Priority: Medium)

5. Expand Test Coverage (4-6 hours)

   • Add tests for __init__.py utility functions
   • Add edge case tests
   • Add integration tests

6. Improve Documentation (3-4 hours)

   • Enhance module docstring
   • Add troubleshooting section to README
   • Standardize docstring format
   • Add DOI links to references

7. Code Quality Improvements (2-3 hours)

   • Consolidate duplicate functions
   • Remove deprecation warnings or update code
   • Improve parameter validation

8. Performance Optimization (2-4 hours)

   • Profile critical functions
   • Add caching where appropriate
   • Optimize NumPy usage

Phase 3: Advanced Features (Priority: Low)

9. API Enhancements (3-5 hours)

   • Consider dataclasses for complex parameters
   • Add convenience functions for common workflows
   • Improve parameter ordering consistency

10. Advanced Testing (2-3 hours)

    • Add property-based tests
    • Add performance benchmarks
    • Add numerical stability tests

11. Documentation Expansion (2-3 hours)

    • Add tutorial examples
    • Create comparison guides (e.g., different LD measures)
    • Add FAQ section

---

9. Detailed Recommendations by File

9.1 __init__.py

Issues:

• Missing exports
• Placeholder implementation
• Inconsistent error handling

Actions:

# Add to __all__:
__all__ = [
    # ... existing exports ...
    "correlation_coefficient",
    "linear_regression",
    "fisher_exact_test",  # or remove if not implemented
    "shannon_entropy",
    "jensen_shannon_divergence",
]

# Fix fisher_exact_test:
def fisher_exact_test(...) -> tuple[float, float]:
    """Calculate Fisher's exact test...
    
    Note: p-value calculation requires scipy.stats.fisher_exact.
    This function returns odds ratio only. For full implementation,
    use scipy.stats.fisher_exact directly.
    """
    # Return odds ratio, raise NotImplementedError for p-value
    # OR implement using scipy with fallback

9.2 coalescent.py

Issues:

• Duplicate watterson_theta functions
• Deprecated parameter order
• Incomplete fay_wu_h implementation

Actions:

• Remove deprecated parameter order support
• Consolidate watterson_theta functions
• Complete fay_wu_h or document limitations

9.3 popgen.py

Issues:

• Duplicate watterson_theta
• Inconsistent error handling (raises vs returns)

Actions:

• Remove duplicate, use coalescent version
• Standardize error handling

9.4 popgen_stats.py

Issues:

• Requires numpy but not clearly documented
• Some functions have placeholder implementations

Actions:

• Document numpy requirement
• Complete placeholder implementations
• Add better fallbacks when scipy unavailable

9.5 README.md

Issues:

• Incorrect module references
• Missing utility functions documentation
• No troubleshooting section

Actions:

• Fix all module references
• Add section on __init__.py utilities
• Add troubleshooting/common errors section
• Add performance notes

---

10. Testing Strategy

10.1 New Tests Needed

1. Unit Tests for __init__.py utilities:

   # tests/math/test_math_utilities.py
   def test_correlation_coefficient():
       # Test perfect correlation
       # Test no correlation
       # Test negative correlation
       # Test edge cases
   
   def test_linear_regression():
       # Test simple line
       # Test no slope
       # Test edge cases
   
   def test_shannon_entropy():
       # Test maximum entropy
       # Test minimum entropy
       # Test normalization
   
   def test_jensen_shannon_divergence():
       # Test identical distributions
       # Test different distributions
       # Test edge cases
   

2. Edge Case Tests:

   # tests/math/test_math_edge_cases.py
   def test_large_sample_sizes():
       # Test with n=1000000
   
   def test_extreme_values():
       # Test with very small/large parameters
   
   def test_numerical_precision():
       # Test floating-point precision limits
   

3. Integration Tests:

   # tests/math/test_math_integration.py
   def test_price_equation_with_popgen():
       # Combine price equation with population genetics
   
   def test_workflow_examples():
       # Test complete analysis workflows
   

10.2 Test Improvements

• Add property-based tests (hypothesis)
• Add performance benchmarks
• Add numerical stability tests
• Improve test documentation

---

11. Documentation Improvements

11.1 Module Docstring Enhancement

"""Mathematical and theoretical biology utilities.

This subpackage provides quantitative biology primitives used across domains:
- Price equation and selection decomposition
- Kin and multilevel selection
- Drift–Diffusion models (DDM)
- Population genetics models
- Coalescent theory
- Epidemiological models
- Linkage disequilibrium
- Quantitative genetics

Submodules:
- coalescent: Coalescent theory and neutrality tests
- price: Price equation and selection metrics
- popgen: Population genetics models
- ld: Linkage disequilibrium
- epidemiology: Disease dynamics models
- dynamics: Population dynamics
- ddm: Decision-making models
- selection: Kin and multilevel selection
- quantgen: Quantitative genetics
- selection_experiments: Natural selection simulations

Utility Functions (in this module):
- correlation_coefficient: Pearson correlation
- linear_regression: Least-squares regression
- shannon_entropy: Information entropy
- jensen_shannon_divergence: Distribution divergence

Dependencies:
- numpy: Required for popgen_stats
- scipy: Optional, for advanced statistics (hardy_weinberg_test, etc.)

Examples:
    >>> from metainformant.math import price_equation, tajimas_D
    >>> # Price equation analysis
    >>> fitness = [1.0, 1.2, 0.9]
    >>> traits = [0.2, 0.4, 0.1]
    >>> cov, trans, total = price_equation(fitness, traits)
"""

11.2 README.md Additions

Add sections:

• Troubleshooting
• Performance Guidelines
• Utility Functions Reference
• Common Workflows
• FAQ

11.3 Function Docstring Standardization

Standard template:

def function_name(param: type) -> return_type:
    """Brief description.
    
    Detailed description of what the function does, including
    mathematical background if relevant.
    
    Args:
        param: Description with constraints/range
    
    Returns:
        Description of return value with units/range
    
    Raises:
        ValueError: When and why
    
    Examples:
        >>> function_name(example_value)
        expected_result
    
    References:
        Author (Year). Title. Journal, volume, pages.
        DOI: 10.xxxx/xxxxx (if available)
    
    Note:
        Any important limitations or assumptions
    """

---

12. Code Quality Metrics

12.1 Current Metrics

• Lines of Code: ~4,500 (excluding tests)
• Functions: ~113
• Classes: 2 (CoalescentSummary, GenerationResult, GenerationsResult)
• Test Files: 22
• Test Coverage: ~85% (estimated)

12.2 Target Metrics

• Test Coverage: >95%
• Documentation Coverage: 100%
• Type Coverage: 100% (already achieved)
• Cyclomatic Complexity: <10 for all functions

---

13. Implementation Priority

Week 1: Critical Fixes

1. Fix placeholder implementations
2. Complete exports
3. Fix documentation errors

Week 2: Test Coverage

4. Add utility function tests
5. Add edge case tests
6. Improve existing tests

Week 3: Documentation

7. Enhance module docs
8. Update README
9. Standardize docstrings

Week 4: Code Quality

10. Consolidate duplicates
11. Standardize error handling
12. Performance optimization

---

14. Success Criteria

Must Have (v1.0)

• All functions exported correctly
• No placeholder implementations
• 100% docstring coverage

• 90% test coverage

• All documentation accurate

Should Have (v1.1)

• Consistent error handling
• Performance optimizations
• Enhanced examples

Nice to Have (v1.2)

• Tutorial guides
• Performance benchmarks
• Advanced workflows

---

15. Conclusion

The math module is well-designed and mathematically rigorous. The main improvements needed are:

1. Completeness: Fix placeholder implementations and missing exports
2. Consistency: Standardize error handling and documentation
3. Testing: Expand test coverage, especially for utilities
4. Documentation: Fix inaccuracies and enhance with examples

With these improvements, the module will be production-ready and maintainable.

---

Reviewer: AI Assistant
Review Date: 2025-01-06
Next Review: After Phase 1 completion