# Design Document: gpa-parallel-njobs

## Overview

**Purpose**: This feature activates the existing `n_jobs` parameter in `GeneralizedProcrustesAnalysis` to enable parallel per-specimen computation during GPA alignment. The `n_jobs` parameter is already declared in the constructor but never wired to `sklearn.utils.parallel.Parallel`/`delayed`.

**Users**: Researchers performing GPA on large specimen collections (hundreds to thousands of specimens) will benefit from multi-core acceleration in both shape and size-and-shape analysis modes, with and without semilandmarks.

**Impact**: Modifies `_transform_shape`, `_transform_size_and_shape`, and `_slide_semilandmarks` in the existing `GeneralizedProcrustesAnalysis` class. No public API changes (the `n_jobs` parameter already exists).

### Goals
- Wire `n_jobs` to `Parallel`/`delayed` for all GPA computation paths (shape, size-and-shape, semilandmark)
- Maintain numerical equivalence regardless of `n_jobs` value
- Follow the established parallelization pattern from `EllipticFourierAnalysis` and `SphericalHarmonicAnalysis`

### Non-Goals
- Adding a `verbose` parameter (out of scope; see `research.md` for rationale)
- Parallelizing across iterations (impossible due to sequential mean-shape dependency)
- Performance benchmarking or auto-tuning of `n_jobs` thresholds

## Architecture

### Existing Architecture Analysis

GPA is iterative: a while-loop alternates between (1) per-specimen alignment to the current mean shape and (2) mean-shape recomputation. Parallelization targets the per-specimen alignment step within each iteration.

Three code paths exist:
1. **Shape mode, no semilandmarks** (`_transform_shape`, line 277): Single list comprehension calling `scipy.spatial.procrustes` per specimen
2. **Shape mode, with semilandmarks** (`_transform_shape`, lines 256-274): Three sequential per-specimen loops (rotate, slide+reproject, center+scale)
3. **Size-and-shape mode** (`_transform_size_and_shape`, line 303): Single list comprehension calling `scipy.linalg.orthogonal_procrustes` per specimen

All paths share the property that per-specimen work within an iteration depends only on the shared read-only mean shape `mu`, making them embarrassingly parallel.

### Architecture Pattern & Boundary Map

```mermaid
graph TB
    subgraph GPA_transform
        FIT[fit_transform called]
        FIT --> BRANCH{scaling?}
        BRANCH -->|True| SHAPE[_transform_shape]
        BRANCH -->|False| SIZESHAPE[_transform_size_and_shape]
        SHAPE --> SEMI{curves?}
        SEMI -->|None| NOSEMI[No-semilandmark path]
        SEMI -->|defined| WITHSEMI[Semilandmark path]
    end

    subgraph Parallel_dispatch
        NOSEMI -->|Parallel| PROCRUSTES[_align_single_shape]
        SIZESHAPE -->|Parallel| ORTHO[_align_single_size_and_shape]
        WITHSEMI -->|Parallel| COMBINED[_iterate_single_semilandmark]
    end
```

**Architecture Integration**:
- Selected pattern: Hybrid — trivial `Parallel` wrapping for simple paths, consolidated per-specimen function for semilandmark path
- Existing patterns preserved: sklearn `Parallel`/`delayed` import pattern, `n_jobs` constructor parameter
- New components: Three private static/helper methods extracted for `delayed` dispatch
- Steering compliance: Follows `tech.md` parallelization convention (`sklearn.utils.parallel.Parallel` and `delayed` with `n_jobs`)

### Technology Stack

| Layer | Choice / Version | Role in Feature | Notes |
|-------|------------------|-----------------|-------|
| Parallelization | `sklearn.utils.parallel.Parallel`, `delayed` | Per-specimen parallel dispatch | Already a dependency (sklearn >= 1.5) |
| Computation | `scipy.spatial.procrustes`, `scipy.linalg.orthogonal_procrustes` | Per-specimen alignment | Existing usage, no change |

No new dependencies introduced.

## System Flows

### Per-Iteration Parallel Dispatch (Semilandmark Path)

```mermaid
sequenceDiagram
    participant TF as _transform_shape
    participant P as Parallel
    participant W1 as Worker 1
    participant W2 as Worker N

    loop while diff_disp > tol
        TF->>P: dispatch _iterate_single_semilandmark for each specimen
        P->>W1: specimen 0 + mu + curves
        P->>W2: specimen N + mu + curves
        W1-->>P: aligned_0, curve_geom_0, disp_0
        W2-->>P: aligned_N, curve_geom_N, disp_N
        P-->>TF: collected results
        TF->>TF: recompute mu from aligned specimens
    end
```

## Requirements Traceability

| Requirement | Summary | Components | Interfaces | Flows |
|-------------|---------|------------|------------|-------|
| 1.1 | Parallel shape alignment with `Parallel`/`delayed` | `_transform_shape`, `_align_single_shape` | `_align_single_shape(mu, x) -> (mu_out, x_aligned, disp)` | No-semilandmark dispatch |
| 1.2 | Sequential when `n_jobs=None` | All paths | — | `Parallel(n_jobs=None)` defaults to sequential |
| 1.3 | Sequential via `Parallel(n_jobs=1)` | All paths | — | Same mechanism |
| 1.4 | Numerical equivalence across `n_jobs` | All paths | — | Tests verify `assert_array_almost_equal` |
| 2.1 | Parallel size-and-shape alignment | `_transform_size_and_shape`, `_align_single_size_and_shape` | `_align_single_size_and_shape(x, mu) -> (x_aligned, disp)` | Size-and-shape dispatch |
| 2.2 | Numerical equivalence in size-and-shape | `_transform_size_and_shape` | — | Tests verify |
| 3.1 | Parallel semilandmark sliding+reproject | `_transform_shape`, `_iterate_single_semilandmark` | `_iterate_single_semilandmark(...)` | Semilandmark dispatch |
| 3.2 | Parallel rotation in semilandmark loop | Consolidated into `_iterate_single_semilandmark` | — | Same function |
| 3.3 | Numerical equivalence with semilandmarks | `_iterate_single_semilandmark` | — | Tests verify |
| 4.1 | Import from `sklearn.utils.parallel` | Module-level import | — | — |
| 4.2 | Pass `n_jobs=self.n_jobs` to `Parallel` | All `Parallel` call sites | — | — |
| 4.3 | `n_jobs=None` defaults to sklearn behavior | All paths | — | — |
| 5.1 | Parametrized `n_jobs` tests (shape, no semilandmarks) | `test_procrustes_analysis.py` | — | — |
| 5.2 | Parametrized `n_jobs` tests (shape, with semilandmarks) | `test_procrustes_analysis.py` | — | — |
| 5.3 | Numerical closeness assertion | `test_procrustes_analysis.py` | — | — |
| 5.4 | Test co-location | `ktch/landmark/tests/test_procrustes_analysis.py` | — | — |

## Components and Interfaces

| Component | Domain/Layer | Intent | Req Coverage | Key Dependencies | Contracts |
|-----------|-------------|--------|--------------|------------------|-----------|
| `_align_single_shape` | landmark / computation | Per-specimen Procrustes alignment (shape, no semilandmarks) | 1.1 | `scipy.spatial.procrustes` (P0) | Service |
| `_align_single_size_and_shape` | landmark / computation | Per-specimen orthogonal rotation (size-and-shape) | 2.1 | `scipy.linalg.orthogonal_procrustes` (P0) | Service |
| `_iterate_single_semilandmark` | landmark / computation | Per-specimen rotate+slide+reproject+center+scale (semilandmark) | 3.1, 3.2 | `_slide_procrustes` / `_slide_bending_energy`, `_reproject_onto_curves` (P0) | Service |
| `_transform_shape` (modified) | landmark / orchestration | Dispatch parallel alignment per iteration | 1.1–1.4, 3.1–3.3 | `Parallel`, `delayed` (P0) | — |
| `_transform_size_and_shape` (modified) | landmark / orchestration | Dispatch parallel alignment per iteration + bug fix | 2.1–2.2 | `Parallel`, `delayed` (P0) | — |

### Landmark / Computation

#### `_align_single_shape` (static method)

| Field | Detail |
|-------|--------|
| Intent | Perform Procrustes alignment of a single specimen to mean shape |
| Requirements | 1.1 |

**Responsibilities & Constraints**
- Wraps `scipy.spatial.procrustes(mu, x)` as a pure function
- Returns aligned specimen and disparity value
- No side effects; suitable for `delayed` dispatch

**Dependencies**
- External: `scipy.spatial.procrustes` — Procrustes alignment (P0)

**Contracts**: Service [x]

##### Service Interface
```python
@staticmethod
def _align_single_shape(
    mu: np.ndarray,  # (n_landmarks, n_dim)
    x: np.ndarray,   # (n_landmarks, n_dim)
) -> tuple[np.ndarray, np.ndarray, float]:
    """Returns (mu_out, x_aligned, disparity)."""
    ...
```
- Preconditions: `mu` and `x` have same shape `(n_landmarks, n_dim)`
- Postconditions: `x_aligned` is optimally rotated/scaled to match `mu`; `disparity` is the sum of squared differences
- Invariants: Pure function, no mutation

**Implementation Notes**
- Directly delegates to `scipy.spatial.procrustes`
- Returns the full 3-tuple from `procrustes` for downstream unpacking

#### `_align_single_size_and_shape` (static method)

| Field | Detail |
|-------|--------|
| Intent | Perform orthogonal rotation of a single specimen to mean shape (no scaling) |
| Requirements | 2.1 |

**Responsibilities & Constraints**
- Wraps `scipy.linalg.orthogonal_procrustes(x, mu)` and applies rotation
- Returns rotated specimen and per-specimen disparity
- No side effects

**Dependencies**
- External: `scipy.linalg.orthogonal_procrustes` — optimal rotation matrix (P0)

**Contracts**: Service [x]

##### Service Interface
```python
@staticmethod
def _align_single_size_and_shape(
    x: np.ndarray,   # (n_landmarks, n_dim)
    mu: np.ndarray,   # (n_landmarks, n_dim)
) -> tuple[np.ndarray, float]:
    """Returns (x_rotated, disparity)."""
    ...
```
- Preconditions: `x` and `mu` have same shape
- Postconditions: `x_rotated = x @ R` where `R` minimizes `||x @ R - mu||`
- Invariants: Pure function

#### `_iterate_single_semilandmark` (method)

| Field | Detail |
|-------|--------|
| Intent | Perform one iteration's per-specimen work: rotate, slide, reproject, center, scale |
| Requirements | 3.1, 3.2 |

**Responsibilities & Constraints**
- Consolidates the 3 per-specimen loops in the semilandmark GPA path into a single function
- Operates on a single specimen against the shared mean shape `mu`
- Returns updated specimen, updated curve geometry, and disparity

**Dependencies**
- Inbound: `_slide_procrustes` or `_slide_bending_energy` — semilandmark sliding (P0)
- Inbound: `_reproject_onto_curves` — curve re-projection (P0)
- External: `scipy.linalg.orthogonal_procrustes` — rotation (P0)

**Contracts**: Service [x]

##### Service Interface
```python
def _iterate_single_semilandmark(
    self,
    x: np.ndarray,            # (n_landmarks, n_dim)
    x_curve_geom: np.ndarray, # (n_landmarks, n_dim)
    mu: np.ndarray,            # (n_landmarks, n_dim)
    curves: np.ndarray,        # (n_sliders, 3)
    slide_func: Callable,      # self._slide_procrustes or self._slide_bending_energy
) -> tuple[np.ndarray, np.ndarray, float]:
    """Returns (x_aligned, x_curve_geom_updated, disparity)."""
    ...
```
- Preconditions: All arrays have compatible shapes; `slide_func` is one of the sliding methods
- Postconditions: `x_aligned` is rotated, slid, reprojected, centered, and unit-scaled; `disparity` is pre-sliding Procrustes distance
- Invariants: Does not modify input arrays; depends only on `mu` (read-only shared state)

**Implementation Notes**
- Steps: (1) `orthogonal_procrustes(x, mu)` → rotate both `x` and `x_curve_geom`, (2) compute disparity, (3) `slide_func(x_rot, mu, curves)`, (4) `_reproject_onto_curves(x_slid, x_cg_rot, curves)`, (5) center and scale both `x` and `x_curve_geom`
- The `slide_func` callable is passed to avoid branching inside the hot loop

### Landmark / Orchestration (Modified Methods)

#### `_transform_shape` (modified)

| Field | Detail |
|-------|--------|
| Intent | Orchestrate iterative GPA with parallel dispatch |
| Requirements | 1.1–1.4, 3.1–3.3 |

**Modification Summary**:
- **No-semilandmark branch**: Replace list comprehension at line 277 with `Parallel(n_jobs=self.n_jobs)(delayed(self._align_single_shape)(mu, x) for x in X_)`
- **Semilandmark branch**: Replace the 3 for-loops (lines 259-274) with single `Parallel(n_jobs=self.n_jobs)(delayed(self._iterate_single_semilandmark)(...) for i in range(n))`
- Collect results and unpack into `X_`, `X_curve_geom`, `total_disp`

#### `_transform_size_and_shape` (modified + bug fix)

| Field | Detail |
|-------|--------|
| Intent | Orchestrate iterative GPA (no scaling) with parallel dispatch |
| Requirements | 2.1–2.2 |

**Modification Summary**:
- **Bug fix**: Replace `self.n_specimen_` / `self.n_landmarks_` with local variables; fix `X` → `X_` in the iteration loop
- **Parallelization**: Replace list comprehension at line 303 with `Parallel(n_jobs=self.n_jobs)(delayed(self._align_single_size_and_shape)(x, mu) for x in X_)`

## Data Models

No data model changes. All inputs and outputs remain `numpy.ndarray`.

## Error Handling

### Error Strategy
No new error categories. The `Parallel` backend handles worker failures internally. If a specimen alignment raises an exception, `Parallel` propagates it to the caller unchanged.

### Existing Validation
- `n_jobs` is validated by sklearn's `Parallel` (accepts `None`, positive int, `-1`, etc.)
- Semilandmark parameters are validated in `_validate_semilandmark_params` (unchanged)

## Testing Strategy

### Unit Tests

Tests are added to `ktch/landmark/tests/test_procrustes_analysis.py`.

| Test | Requirements | Description |
|------|-------------|-------------|
| `test_gpa_shape_parallel` | 1.1, 1.2, 1.3, 1.4, 5.1, 5.3 | Parametrized `n_jobs=[None, 1, 3]`; verify `fit_transform` output and `mu_` are `assert_array_almost_equal` across all values |
| `test_gpa_semilandmark_parallel_bending` | 3.1, 3.2, 3.3, 5.2, 5.3 | Parametrized `n_jobs=[None, 1, 3]` with `sliding_criterion='bending_energy'`; verify numerical equivalence |
| `test_gpa_semilandmark_parallel_procrustes` | 3.1, 3.2, 3.3, 5.2, 5.3 | Same as above with `sliding_criterion='procrustes_distance'` |
| `test_gpa_size_and_shape_parallel` | 2.1, 2.2, 5.3 | Parametrized `n_jobs=[None, 1, 3]` for `scaling=False`; verify equivalence |

### Test Pattern
Following the EFA test pattern (`test_elliptic_fourier_analysis.py:88-106`):
```python
@pytest.mark.parametrize("n_jobs", [None, 1, 3])
def test_gpa_shape_parallel(n_jobs):
    gpa = GeneralizedProcrustesAnalysis(n_dim=2, n_jobs=n_jobs)
    X_aligned = gpa.fit_transform(X)
    # Compare against baseline (n_jobs=None)
    ...
    assert_array_almost_equal(X_aligned, X_baseline)
```

## Performance & Scalability

- **Overhead**: `Parallel` is invoked once per iteration (not per specimen). With typical GPA convergence in 5–15 iterations, overhead is negligible for datasets with >50 specimens.
- **Scaling**: Linear speedup expected for embarrassingly parallel per-specimen operations. Semilandmark sliding (involving matrix operations) benefits most from parallelization.
- **Default behavior**: `n_jobs=None` preserves current sequential performance with zero overhead (no `Parallel` instantiation path needed — `Parallel(n_jobs=None)` in sklearn runs sequentially in the calling process).
