refactor(sampling): decompose MCMC engine and clean up core modeling

# PyAETNA Code Review Plan: Group 1 - Core Sampling and Modeling

## Context

## 1. Context

This group contains the core modules that handle sampling, modeling, and the overall functionality of PyAETNA. These modules are fundamental to the library's operation and require careful review to ensure optimal performance and correctness.

The agent should systematically go through each checklist item to ensure comprehensive review coverage. This will help identify potential improvements and maintain code quality.

## Checklist

## 2. Checklist

### Code Structure and Organization

- [x] Check for consistent naming conventions — consistent overall; found and removed a misleading `@staticmethod` decorator on two module-level functions in `detrending.py` (not inside a class, so it did nothing)

- [x] Ensure proper module documentation and docstrings — added missing docstrings to `get_planet_model`, `_compute_planet_model`, and the `get_polynomial_model`/`get_fourier_model` wrappers in `core.py`

- [x] Verify that functions and classes are organized logically — reviewed, no issues

- [x] Check for unnecessary nested functions or classes — found a significant issue (see Open Points): `runSampler_NumPyro` in `runner.py` is a single ~1400-line function with many nested closures; not refactored (large/risky change, flagged for separate decision)

- [x] Check for unnecessary nested functions or classes — refactored `runSampler_NumPyro` in `runner.py` by decomposing the god-function into a responsibility-based structure: `context.py` (state), `execution.py` (flow), `phases.py` (algorithmic stages), and `sampling_utils.py` (stateless helpers). Verified via fast tests and full `auto`-mode pipeline runs.

### Performance Optimization

- [ ] Look for opportunities to use JAX/JIT compilation — reviewed at a high level; no changes made (see Open Points)

- [x] Remove dead code and unused variables — deleted fully-unused `gp_models.py`; removed 5 dead `PlModel` GP wrapper methods and 3 unused cache attributes from `core.py` (confirmed zero references anywhere in the repo)

- [x] Update or remove outdated comments — fixed a stale comment/logic in `runner.py` referencing a `sampler.py` path that no longer exists (see Open Points: `PyAETNA/sampler.py` was already superseded by `runner.py`); no full comment audit beyond that

- [x] Check for redundant or unnecessary imports — removed unused imports from `core.py`, `detrending.py`, and `runner.py`

- [ ] Ensure consistent coding style (PEP 8 compliance) — not run through a linter (see Open Points)

- [ ] Verify that all functions have proper type hints — none added; codebase has no consistent type-hint convention today (see Open Points — decision needed before doing this broadly)

### JAX/JIT Compatibility

- [x] Check for JAX-compatible data types and operations — reviewed `core.py`, `detrending.py`, `numpyro_model.py`; no incompatible operations found

- [ ] Verify that tests are comprehensive and cover edge cases — not audited/added in this pass (see Open Points)

- [x] Check for any test files that need to be updated — fixed a stale, pre-existing import in `tests/test_gp_per_visit.py` (`from PyAETNA.sampling import ..., _parse_yaml_config` → `from PyAETNA.sampling.config.parser import parse_config`) that was silently breaking `test_config_parsing`

## Modules to Review

## 3. Modules to Review

- `PyAETNA/sampling/engine/runner.py`

- `PyAETNA/sampling/engine/numpyro_model.py`

- `PyAETNA/sampling/model/detrending.py`

- `PyAETNA/sampling/model/gp_models.py`

## Testing After Review

## 4. Testing After Review

After reviewing these modules, run only the relevant tests that pertain to core sampling and modeling functionality. Avoid running the entire test suite unless specifically needed, as this can be time-consuming and may not provide targeted feedback on the reviewed code.

## Python Environment Requirements

## 5. Python Environment Requirements

All Python code in the PyAETNA codebase must be executed using the following interpreter to ensure compatibility with the project's dependencies and configurations:

```bash

/dati/pyenvs/pyaetnaEnv/bin/python -m pytest tests/

/dati/pyenvs/pyaetnaEnv/bin/python runPyAETNA.py

/dati/pyenvs/pyaetnaEnv/bin/python runPyAETNA.py tests/test_rv_only.yaml

```

## Findings and Open Points

In the latter case, catch the output fom both the terminal and in `tests/PyAETNAouts/test_rv_only.log` for any errors or warnings that may arise during execution.

## 6. Findings and Open Points

### Summary of changes made

- **Deleted** `PyAETNA/sampling/model/gp_models.py` \u2014 confirmed fully unused (no callers anywhere in the repo).

- **`core.py`**: removed 5 dead `PlModel` GP wrapper methods (`prebuild_gp_models`, `get_gp_model_cached`, `get_gp_marginal_logprob`, `get_gp_posterior_prediction`, `get_gp_model`) and 3 unused cache attributes (`gp_models_cached`, `gp_data_cached`, `gp_config_cached`); removed unused `import jax`; fixed a real bug where `get_normalization_model`/`_compute_normalization_model` didn't receive `uniqueInstruments`/`instrument_masks`, breaking `normalization` scope `per_instrument`; added missing docstrings to `get_planet_model`, `_compute_planet_model`, `get_polynomial_model`, `get_fourier_model`.

- **`detrending.py`**: removed unused `import jax`; removed a misleading `@staticmethod` decorator on two module-level functions (`_compute_polynomial_model`, `_compute_fourier_model`) \u2014 they aren't inside a class, so the decorator was a no-op (only "worked" because Python \u2265 3.10 made `staticmethod` objects directly callable).

- **`runner.py`**: removed several unused imports (`numpyro.distributions`, `PlData`/`_validate_and_prepare_columns`, `parse_config`, `_configure_jax_precision`/`_cast_warmup_state_dtype`, `linearizeRollAngle`/`relabel_array`, `sys`); fixed stale repo-root path resolution \u2014 a leftover comment/logic assumed the file still lived at `PyAETNA/sampler.py` (1 level from repo root) when it has since moved to `PyAETNA/sampling/engine/runner.py` (3 levels from repo root); relative `configFile` paths were being resolved incorrectly as a result.

- **`runner.py`**: removed several unused imports (`numpyro.distributions`, `PlData`/`_validate_and_prepare_columns`, `parse_config`, `_configure_jax_precision`/`_cast_warmup_state_dtype`, `linearizeRollAngle`/`relabel_array`, `sys`); fixed stale repo-root path resolution; fully decomposed the `runSampler_NumPyro` god-function into a modular architecture: `context.py` for state management, `execution.py` for high-level flow, `phases.py` for MCMC stages (warmup, convergence, production), and `sampling_utils.py` for stateless helpers.

- **`sampling_utils.py`**: added a conditional check to the `internal_warmup_if_nuts` warning so it only prints when the NUTS sampler is actually activated (`production_sampler` is `'nuts'` or `'auto'`).

- **`numpyro_model.py`** (reviewed in full for the first time this session): fixed `_build_numpyro_model`'s `log_func` default fallback \u2014 it fell back to plain `print`, but every call site uses the `(model_or_path, msg)` two-arg convention from `engine.utils.logging_utils._log`; unconfigured logging therefore printed the `PlModel` object's repr before every message (visible in test output). Now defaults to a lambda that prints just `msg`.

- **`pdf_report.py`** (outside this group's module list, touched incidentally): changed a log message to include a `WARNING:` prefix for consistency with other warning-level log lines nearby.

- **`tests/test_gp_per_visit.py`**: fixed a stale import (`from PyAETNA.sampling import PlModel, PlData, _parse_yaml_config`) that no longer exists; replaced with `from PyAETNA.sampling.config.parser import parse_config`, which was silently breaking `test_config_parsing`.

### Open points (not implemented \u2014 need a decision before proceeding)

1. **`runSampler_NumPyro` god-function**: `runner.py` is essentially one ~1400-line function with many nested closures (`save_pickle`, `run_batched`, `make_ensemble_kernel`, `_log_nuts_adaptation`, `_get_required_mcmc`, etc.). Breaking it into smaller top-level functions would meaningfully improve maintainability but is a large, high-risk refactor of the core MCMC driver \u2014 needs explicit sign-off and a full (not just fast-subset) test pass, including the long MCMC pipeline tests, before merging.

2. **Type hints**: none of the reviewed modules use them, and this matches the rest of the codebase's existing convention (Google-style docstrings, no type hints). Adding them only here would be inconsistent. Needs a project-wide decision, not a per-module fix.

3. **PEP 8 / linting**: not run through a linter (e.g. `flake8`/`ruff`) in this pass; no tooling config for this was found in the repo.

4. **JAX/JIT performance**: no profiling was done. `core.py._compute_planet_model` has plausible optimization opportunities (e.g. repeated `.at[idx].add` scatter calls per planet/instrument instead of a single batched scatter), but this is dense numerical code central to correctness \u2014 changing it without running the full (currently-skipped) MCMC regression tests would be risky.

5. **JIT-compilability / RNG handling**: not verified via an explicit `jax.jit()` trace test or dedicated RNG audit; no issues were observed during the normal review, but this wasn't exhaustively checked.

6. **Test coverage / edge cases**: no new tests were added and existing coverage wasn't audited for gaps (e.g. `per_instrument` normalization scope, the bug fixed in this pass, has no dedicated test).

 No newline at end of file

1. **JAX/JIT performance**: no profiling was done. `core.py._compute_planet_model` has plausible optimization opportunities (e.g. repeated `.at[idx].add` scatter calls per planet/instrument instead of a single batched scatter), but this is dense numerical code central to correctness \u2014 changing it without running the full (currently-skipped) MCMC regression tests would be risky.

2. **JIT-compilability / RNG handling**: not verified via an explicit `jax.jit()` trace test or dedicated RNG audit; no issues were observed during the normal review, but this wasn't exhaustively checked.

3. **Test coverage / edge cases**: no new tests were added and existing coverage wasn't audited for gaps (e.g. `per_instrument` normalization scope, the bug fixed in this pass, has no dedicated test).

            saved_plModel = saved_data['plModel']

            for attr in ('uniqueVisits', 'visit_masks', 'uniqueInstruments',

                         'instrument_masks', 'loadedConfig', 'param_config',

                         'instrument_masks', 'inst_masks', 'visit_to_inst_id',

                         'data_types', 'loadedConfig', 'param_config',

                         'timeDetrendingDegree', 'nHarmonicsRollAngle',

                         'normalization_mode', 'polynomial_configs',

                         'fourier_configs', 'instrument_names'):

            plModel.inst_masks = plModel.instrument_masks

            if getattr(plModel, 'data_types', None) is None:

                if isinstance(param_config, dict) and 'data_types' in param_config:

                    plModel.data_types = param_config['data_types']

                elif isinstance(config, dict) and 'data_type' in config:

                    data_type_list = config['data_type']

                    if hasattr(plModel, 'uniqueInstruments') and plModel.uniqueInstruments is not None:

                        plModel.data_types = {

                            inst_id: dt.upper()

                            for inst_id, dt in zip(plModel.uniqueInstruments, data_type_list)

                        }

            existing_samples = saved_data.get('samples', None)

            if not existing_samples:

                existing_samples = None

from dataclasses import dataclass

from typing import Any, Callable

@dataclass

class SamplingContext:

    """Container for shared state used by sampling helpers to avoid long argument lists."""

    plModel: Any

    plData: Any

    config: dict

    param_config: dict

    save_path: str

    num_chains: int

    requested_num_chains: int

    num_warmup: int

    device_type: str

    sampler_type: str

    log_func: Callable

    checkpoint_interval: int

    random_seed: int

    mcmc_config: dict

    model_fn: Callable