first implementation fo refactoring of group1 modules. code working, tests...

# PyAETNA Code Review Plan

This plan outlines a systematic approach to reviewing the PyAETNA codebase for performance optimization, JAX/JIT compatibility, and code quality.

## 1. Group the Modules by Functionality

### Group 1: Core Sampling and Modeling

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

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

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

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

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

### Group 2: Kernel and GP-related Modules

- `PyAETNA/gp.py`

- `PyAETNA/kernels/factory.py`

- `PyAETNA/kernels/spectral_mixture.py`

### Group 3: Data Loading and Processing

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

- `PyAETNA/sampling/data.py`

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

### Group 4: Reporting and Visualization

- `PyAETNA/reporting/confidence_bands.py`

- `PyAETNA/reporting/lightcurve_plots.py`

- `PyAETNA/reporting/model_components.py`

- `PyAETNA/reporting/pdf_report.py`

- `PyAETNA/reporting/periodogram_plots.py`

- `PyAETNA/reporting/plot_models.py`

- `PyAETNA/reporting/plot_styling.py`

- `PyAETNA/reporting/panel_detrended.py`

- `PyAETNA/reporting/panel_nontime_residuals.py`

- `PyAETNA/reporting/panel_overlay.py`

- `PyAETNA/reporting/panel_raw.py`

- `PyAETNA/reporting/panel_residuals.py`

- `PyAETNA/reporting/title_formatting.py`

- `PyAETNA/reporting/ttv_plots.py`

### Group 5: Configuration and Parsing

- `PyAETNA/sampling/config/parser/__init__.py`

- `PyAETNA/sampling/config/parser/detrending.py`

- `PyAETNA/sampling/config/parser/gp.py`

- `PyAETNA/sampling/config/parser/planets.py`

- `PyAETNA/sampling/config/parser/prior_refs.py`

- `PyAETNA/sampling/config/parser/scoping.py`

- `PyAETNA/sampling/config/parser/star.py`

- `PyAETNA/sampling/config/priors.py`

### Group 6: Utilities and Helpers

- `PyAETNA/sampling/utils.py`

- `PyAETNA/sampling/precision.py`

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

- `PyAETNA/sampling/engine/utils/logging_utils.py`

### Group 7: Logging and Error Handling

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

- `PyAETNA/reporting/logger.py`

### Group 8: Tests and Examples

- `tests/test_gap_detection.py`

- `tests/test_gap_plots.py`

- `tests/test_gp_per_visit.py`

- `tests/test_lc_gp_rv.yaml`

- `tests/test_lc_only.yaml`

- `tests/test_phase_curve.yaml`

- `tests/test_prior_refs.py`

- `tests/test_rv_only.yaml`

- `tests/test_spectral_mixture.py`

- `tests/test_spectral_mixture.yaml`

- `tests/test_spectral_mixture_mcmc.py`

## 2. Checklist for Each Module

For each module, perform the following checks:

### Code Structure and Organization

- [ ] Check for consistent naming conventions

- [ ] Ensure proper module documentation and docstrings

- [ ] Verify that functions and classes are organized logically

- [ ] Check for unnecessary nested functions or classes

### Performance Optimization

- [ ] Look for opportunities to use JAX/JIT compilation

- [ ] Identify and optimize performance bottlenecks

- [ ] Check for inefficient data structures or algorithms

- [ ] Ensure that vectorized operations are used where possible

### Code Quality

- [ ] Remove dead code and unused variables

- [ ] Update or remove outdated comments

- [ ] Check for redundant or unnecessary imports

- [ ] Ensure consistent coding style (PEP 8 compliance)

- [ ] Verify that all functions have proper type hints

### JAX/JIT Compatibility

- [ ] Check for JAX-compatible data types and operations

- [ ] Ensure that all functions can be JIT-compiled

- [ ] Verify that random number generation is handled correctly for JAX

- [ ] Check for any Python features not supported by JAX

### Error Handling and Logging

- [ ] Ensure proper error handling and logging

- [ ] Check for consistent logging levels and formats

- [ ] Verify that all exceptions are caught and handled appropriately

### Testing

- [ ] Ensure that all functions are covered by tests

- [ ] Verify that tests are comprehensive and cover edge cases

- [ ] Check for any test files that need to be updated

## 3. Review Process

1. **Start with Core Modules**: Begin with Group 1 (Core Sampling and Modeling) as these are fundamental to the library's functionality.

2. **Move to Dependencies**: Review Group 2 (Kernel and GP-related Modules) next, as they are closely tied to the core functionality.

3. **Data Processing**: Then move to Group 3 (Data Loading and Processing) to ensure data handling is efficient and correct.

4. **Reporting**: Review Group 4 (Reporting and Visualization) to ensure all visualizations are accurate and performant.

5. **Configuration**: Check Group 5 (Configuration and Parsing) to ensure all configuration options are properly handled.

6. **Utilities**: Review Group 6 (Utilities and Helpers) for any general-purpose functions that might need optimization.

7. **Logging**: Check Group 7 (Logging and Error Handling) to ensure consistent and helpful logging.

8. **Tests**: Finally, review Group 8 (Tests and Examples) to ensure all functionality is properly tested.

**Recommendation**: For each group, follow the checklist provided in the corresponding group-specific plan file. The agent should systematically go through each checklist item to ensure comprehensive review coverage.

## 4. Testing After Each Group

After reviewing each group, run only the relevant tests that pertain to the changes made. 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:

- For Groups 1-3: Run core functionality tests only

- For Group 4: Run reporting and visualization tests only

- For Group 5: Run configuration and parsing tests only

- For Groups 6-7: Run utility and logging tests only

- For Group 8: Run the specific tests related to the reviewed examples

## 5. Final Review

After all groups have been reviewed and tested:

- [ ] Perform a final code review for any remaining issues

- [ ] Run a full test suite to ensure everything works as expected

- [ ] Check for any performance regressions

This plan should help you systematically review the codebase while maintaining functionality and improving performance.

## 2. 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

```

This interpreter has been configured with the specific Python version and dependencies required for PyAETNA. Using any other interpreter may result in compatibility issues or missing dependencies.

### Running Tests and Scripts

When running tests or scripts, always use the full path to the interpreter:

```bash

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

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

```

 No newline at end of file

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

## 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.

### Criticalities

1. **Performance**: These modules handle computationally intensive tasks like sampling and model fitting. Any inefficiencies here can significantly impact overall performance.

2. **JAX/JIT Compatibility**: Ensuring that these modules are compatible with JAX/JIT compilation is crucial for achieving optimal performance.

3. **Code Quality**: Maintaining clean, well-documented, and well-structured code is essential for future maintainability.

### Suggestions for Improvements

1. **JAX/JIT Optimization**: Look for opportunities to use JAX/JIT compilation to speed up computations.

2. **Vectorization**: Ensure that operations are vectorized where possible to take advantage of JAX's capabilities.

3. **Code Refactoring**: Refactor any complex or poorly structured code to improve readability and maintainability.

4. **Documentation**: Ensure that all functions and classes have proper docstrings and documentation.

### Recommendation

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

### 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)

### Performance Optimization

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

- [ ] Identify and optimize performance bottlenecks — not profiled; no changes made (see Open Points)

- [ ] Check for inefficient data structures or algorithms — not exhaustively audited (see Open Points)

- [x] Ensure that vectorized operations are used where possible — reviewed `core.py`/`detrending.py`; per-visit/per-instrument Python loops are intentional and correct (static, config-determined counts, unrolled at JAX trace time)

### Code Quality

- [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

- [ ] Ensure that all functions can be JIT-compiled — not verified via an explicit `jax.jit()` trace test (see Open Points)

- [ ] Verify that random number generation is handled correctly for JAX — not specifically audited in this pass

- [x] Check for any Python features not supported by JAX — none observed in the reviewed modules

### Error Handling and Logging

- [x] Ensure proper error handling and logging — reviewed all `try/except` blocks in `runner.py`; all degrade gracefully and log the caught exception (no silent swallowing, no bare `except:`)

- [x] Check for consistent logging levels and formats — fixed a real bug: `_build_numpyro_model`'s `log_func` defaulted to plain `print`, but call sites use the `(model_or_path, msg)` two-arg convention from `logging_utils._log`, so unconfigured logging printed the `PlModel` object's repr before every message. Fixed to default to printing just `msg`.

- [x] Verify that all exceptions are caught and handled appropriately — reviewed, appropriate

### Testing

- [x] Ensure that all functions are covered by tests — ran the fast, relevant tests with `/dati/pyenvs/pyaetnaEnv/bin/python`: `test_behavioral_ties.py`, `test_eiv.py`, `test_gp_per_visit.py` (all pass); long MCMC-pipeline tests (`test_rm_effect.py`, `test_spectral_mixture.py`, `test_ttv.py`) were intentionally skipped per instructions

- [ ] 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

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

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

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

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

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

## 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

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

```

This interpreter has been configured with the specific Python version and dependencies required for PyAETNA. Using any other interpreter may result in compatibility issues or missing dependencies.

### Running Tests and Scripts

When running tests or scripts, always use the full path to the interpreter:

```bash

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

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

```

## 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.

- **`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`.

### Verification performed

- Import/instantiation smoke test of `core.PlModel`, `detrending`, and `runner` via `/dati/pyenvs/pyaetnaEnv/bin/python`; confirmed removed attributes/methods are gone and the repo-root path fix resolves correctly.

- Ran `test_behavioral_ties.py` (2/2 pass), `test_eiv.py` (5/5 pass), `test_gp_per_visit.py` (5/5 pass) with `/dati/pyenvs/pyaetnaEnv/bin/python`.

- Skipped `test_rm_effect.py`, `test_spectral_mixture.py`, `test_ttv.py` \u2014 these run full MCMC pipelines and were excluded as "long tests" per instructions.

- `git grep` confirmed zero remaining references to all removed dead-code symbols anywhere in the repo (source or tests) before deleting them.

### 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

# PyAETNA Code Review Plan: Group 2 - Kernel and GP-related Modules

## Context

This group contains modules that handle Gaussian Processes (GPs) and kernel functions, which are essential for modeling and prediction in PyAETNA. These modules are closely tied to the core functionality and require careful review to ensure correctness and performance.

### Criticalities

1. **Mathematical Correctness**: Ensuring that the GP and kernel implementations are mathematically correct is crucial for accurate modeling.

2. **Performance**: These modules handle computationally intensive tasks like kernel matrix computations. Any inefficiencies here can significantly impact overall performance.

3. **JAX/JIT Compatibility**: Ensuring that these modules are compatible with JAX/JIT compilation is crucial for achieving optimal performance.

### Suggestions for Improvements

1. **JAX/JIT Optimization**: Look for opportunities to use JAX/JIT compilation to speed up kernel matrix computations.

2. **Vectorization**: Ensure that operations are vectorized where possible to take advantage of JAX's capabilities.

3. **Code Refactoring**: Refactor any complex or poorly structured code to improve readability and maintainability.

4. **Documentation**: Ensure that all functions and classes have proper docstrings and documentation.

### Recommendation

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

### Code Structure and Organization

- [ ] Check for consistent naming conventions

- [ ] Ensure proper module documentation and docstrings

- [ ] Verify that functions and classes are organized logically

- [ ] Check for unnecessary nested functions or classes

### Performance Optimization

- [ ] Look for opportunities to use JAX/JIT compilation

- [ ] Identify and optimize performance bottlenecks

- [ ] Check for inefficient data structures or algorithms

- [ ] Ensure that vectorized operations are used where possible

### Code Quality

- [ ] Remove dead code and unused variables

- [ ] Update or remove outdated comments

- [ ] Check for redundant or unnecessary imports

- [ ] Ensure consistent coding style (PEP 8 compliance)

- [ ] Verify that all functions have proper type hints

### JAX/JIT Compatibility

- [ ] Check for JAX-compatible data types and operations

- [ ] Ensure that all functions can be JIT-compiled

- [ ] Verify that random number generation is handled correctly for JAX

- [ ] Check for any Python features not supported by JAX

### Error Handling and Logging

- [ ] Ensure proper error handling and logging

- [ ] Check for consistent logging levels and formats

- [ ] Verify that all exceptions are caught and handled appropriately

### Testing

- [ ] Ensure that all functions are covered by tests

- [ ] Verify that tests are comprehensive and cover edge cases

- [ ] Check for any test files that need to be updated

## Modules to Review

- `PyAETNA/gp.py`

- `PyAETNA/kernels/factory.py`

- `PyAETNA/kernels/spectral_mixture.py`

## Testing After Review

After reviewing these modules, run only the relevant tests that pertain to kernel and GP 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.

## 2. 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

```

This interpreter has been configured with the specific Python version and dependencies required for PyAETNA. Using any other interpreter may result in compatibility issues or missing dependencies.

### Running Tests and Scripts

When running tests or scripts, always use the full path to the interpreter:

```bash

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

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

```

 No newline at end of file