Refactor optim optimizers into readable quasi-private helpers

[kalebphipps]

Here is the first true Claude contribution to ARTIST — no new physics, no new bugs (hopefully), just four 250–530-line methods finally learning to use paragraphs. 🧹🤖

What & why

The optimizers/reconstructors in artist/optim had grown four very long methods (250–530 lines each) that packed setup, data loading, the epoch loop, constraints, and distributed synchronization into a single body — and the two kinematics methods were ~90% copy-paste of each other. This PR breaks each of them into a short orchestrator that reads as a sequence of clearly-named quasi-private (_) helper methods.

This is a pure extract-method refactor. The public API and numerical behaviour of every method are unchanged. Nothing about how you call these classes changes.

Scope

Method	Before → after (orchestrator)
KinematicsReconstructor._reconstruct_kinematics_parameters_with_raytracing	312 → 178
KinematicsReconstructor._reconstruct_kinematics_parameters_with_alignment	258 → ~120
SurfaceReconstructor.reconstruct_surfaces	534 → 328
AimPointOptimizer.optimize	519 → 249

Nothing outside artist/optim/ is touched. loss.py, regularizers.py, and training.py are untouched.

Design decisions

• Per-class private methods only. No new module-level functions, no shared base class, no new imports. This keeps each file self-contained and avoids any coupling / circular-import risk. The (small) cost is that _initialize_reconstruction_bookkeeping and _parse_group_calibration_data are duplicated between KinematicsReconstructor and SurfaceReconstructor — a deliberate, low tax for zero coupling.
• The big duplication is gone. The two ~90%-identical kinematics methods now share helpers (_initialize_reconstruction_bookkeeping, _parse_group_calibration_data, _setup_optimizer_scheduler_early_stopping, _synchronize_*) since they live in the same class.
• Control flow stays visible in the orchestrator. The while loop, break on early stopping, epoch += 1, per-epoch history .append(...), and the loss composition were intentionally not extracted — only cohesive blocks (setup, forward pass, constraints, distributed sync) moved into helpers.
• Stateful semantics are preserved explicitly. Anything that must persist across epochs — the epoch == 0 reference capture and the Augmented-Lagrangian lambda_* multipliers — is round-tripped in/out of the helper rather than hidden inside it. See AimPointOptimizer._compute_kl_constraints and SurfaceReconstructor._compute_flux_integral_constraint.
• Every helper carries a NumPy-style docstring + full type hints, matching the existing _validate / _plot_fluxes convention. This is why the line count grows despite the logic shrinking — the added lines are docstrings and explicit signatures, not new behaviour.

What reviewers should focus on

Because this is behaviour-preserving by construction, the useful review question is "is each helper a faithful move of the original block?" rather than "is the logic correct?". Concretely:

1. AimPointOptimizer._compute_kl_constraints — the trickiest extraction. Check that the three constraints, the loss = flux_loss + ... composition, and the three lambda_* updates under torch.no_grad() are byte-for-byte the original, and that the epoch == 0 references + multipliers correctly round-trip (in as params, out in the return tuple) so state persists across epochs. Also confirm the pre-existing behaviour where loss is only reassigned inside the KL branch is preserved.
2. SurfaceReconstructor._compute_flux_integral_constraint / _compute_regularization_terms — confirm the epoch == 0 flux_integrals_reference capture stays in the orchestrator (guarding the call), and that alpha/beta dynamic balancing and the epsilon handling are unchanged.
3. _synchronize_* helpers (all three classes) — these wrap the is_nested / is_distributed branches. They're no-ops in the single-rank test path, so they're the least test-covered lines; a careful read of the broadcast/all-reduce/all-gather calls is worthwhile.
4. The two kinematics methods — verify the shared helpers are genuinely identical for both, and that the only per-method differences remain inline: nan_to_num_ (alignment) vs _synchronize_gradients_nested_ddp (raytracing), and the _validate reduction (torch.mean vs torch.median).
5. _predict_flux (surface) / _compute_raytracing_loss (kinematics) — confirm the ray-tracer construction args, ordering, and the sample_indices / local_indices derivations match the originals exactly.

A helpful way to review: read each orchestrator top-to-bottom first (it should now tell a clean story), then spot-check that each self._helper(...) call receives/returns exactly what the inlined code used to.

Verification

• Full test suite: 407 passed. The kinematics and surface reconstructor tests do real golden-file comparisons (tests/data/expected_test_data.pt) with tolerances, so their passing genuinely validates behaviour-equivalence.
• Coverage held at 95% (per-file actually nudged up: aim-point 94→95%, kinematics 92→93%). The only uncovered lines are the distributed/nested-DDP branches, which were already uncovered on main.
• ruff check + ruff format clean; pre-commit mypy passes with no new errors versus main.

Non-goals / follow-ups

• Other long functions exist outside optim (heliostat_field.from_hdf5 ~355 lines, heliostat_ray_tracer.trace_rays ~289) — intentionally out of scope here.
• Unrelated observation while testing: tests/optim/test_aim_point_optimizer.py overwrites its own key in expected_test_data.pt and then asserts against what it just wrote, so it always passes and dirties the working tree on every run. Not touched here, but worth a separate look.

🤖 Generated with Claude Code

[github-actions]

Name	Stmts	Miss	Cover	Missing
artist\__init__.py	3	0	100%
artist\field\__init__.py	14	0	100%
artist\field\actuators.py	16	0	100%
artist\field\actuators_ideal.py	9	0	100%
artist\field\actuators_linear.py	57	0	100%
artist\field\heliostat_field.py	116	0	100%
artist\field\heliostat_group.py	52	0	100%
artist\field\heliostat_group_rigid_body.py	23	0	100%
artist\field\kinematics.py	12	0	100%
artist\field\kinematics_rigid_body.py	98	0	100%
artist\field\solar_tower.py	48	0	100%
artist\field\surface.py	22	0	100%
artist\field\tower_target_areas.py	15	0	100%
artist\field\tower_target_areas_cylindrical.py	40	0	100%
artist\field\tower_target_areas_planar.py	32	0	100%
artist\flux\__init__.py	2	0	100%
artist\flux\bitmap.py	65	1	98%	116
artist\geometry\__init__.py	4	0	100%
artist\geometry\coordinates.py	95	1	99%	293
artist\geometry\rotations.py	31	1	97%	112
artist\geometry\transforms.py	97	1	99%	346
artist\io\__init__.py	3	0	100%
artist\io\calibration_parser.py	42	2	95%	197-202
artist\io\h5_scenario_parser.py	149	5	97%	550, 554, 560, 565, 570
artist\io\paint_calibration_parser.py	66	0	100%
artist\io\paint_scenario_parser.py	193	0	100%
artist\io\stral_scenario_parser.py	42	0	100%
artist\nurbs\__init__.py	3	0	100%
artist\nurbs\surfaces.py	156	20	87%	210-241, 395, 399, 401-402, 553-554, 557-558
artist\nurbs\utils.py	24	0	100%
artist\optim\__init__.py	7	0	100%
artist\optim\aim_point_optimizer.py	186	10	95%	516, 695-702, 713-722
artist\optim\kinematics_reconstructor.py	207	15	93%	640-648, 677-700, 832, 1001
artist\optim\loss.py	85	0	100%
artist\optim\regularizers.py	24	0	100%
artist\optim\surface_reconstructor.py	212	12	94%	769-777, 817-835
artist\optim\training.py	67	1	99%	181
artist\raytracing\__init__.py	4	0	100%
artist\raytracing\blocking.py	283	9	97%	378, 380, 552-555, 592-593, 722, 737, 975
artist\raytracing\geometry.py	91	0	100%
artist\raytracing\heliostat_ray_tracer.py	142	1	99%	176
artist\raytracing\sampling.py	29	0	100%
artist\scenario\__init__.py	4	0	100%
artist\scenario\h5_scenario_generator.py	83	30	64%	117-141, 153-161, 185, 192-197, 215-218
artist\scenario\scenario.py	94	2	98%	215, 434
artist\scenario\surface_generator.py	89	0	100%
artist\scene\__init__.py	5	0	100%
artist\scene\light_source.py	11	0	100%
artist\scene\light_source_array.py	28	0	100%
artist\scene\rays.py	7	0	100%
artist\scene\sun.py	43	0	100%
artist\util\__init__.py	56	18	68%	101-135
artist\util\config.py	203	0	100%
artist\util\constants.py	305	0	100%
artist\util\env.py	110	23	79%	89-91, 186-228, 308
artist\util\indices.py	324	0	100%
artist\util\type_registry.py	11	0	100%
TOTAL	4239	152	96%

[codecov]

Codecov Report

❌ Patch coverage is 85.10638% with 35 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.41%. Comparing base (63390f3) to head (fcfec2c).

Files with missing lines	Patch %	Lines
artist/optim/kinematics_reconstructor.py	81.94%	13 Missing ⚠️
artist/optim/surface_reconstructor.py	84.41%	12 Missing ⚠️
artist/optim/aim_point_optimizer.py	88.37%	10 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #220      +/-   ##
==========================================
+ Coverage   96.38%   96.41%   +0.02%     
==========================================
  Files          57       57              
  Lines        4208     4239      +31     
==========================================
+ Hits         4056     4087      +31     
  Misses        152      152              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[MarleneBusch]

I added a fix in ‎artist/optim/training.py‎. In the last PR Marie moved some tensors which were only used for indexing to cpu, but missed one which broke my computations on gpu. Its fixed now.

[mcw92]

LGTM. I checked your comment regarding the :meth: directive @MarleneBusch. This turns out fine in the API references with a link to the respective method.
I did not check everything line by line for equivalence with the code before. However, this is ok for me as the tests pass, Claude was told to only do a refactor, and @MarleneBusch already checked everything carefully.
I think the optimize methods are now way more readable than before (even though there might still be room for improvement somewhere, as always). I will merge this now.