refactor(fk_steering): extract per-particle guided step into a helper

[manzuoni-astera]

Summary

Fixes #77.

FKSteering._run_step was a ~125-line method that mixed two concerns: the unguided batched fast path and the guided per-particle loop (with its 4D restacking and SamplerStepOutput reassembly). Extracted the guided branch into _run_guided_step_per_particle, so _run_step now reads as a small dispatch:

coords_4d = coords.reshape(num_particles, self.ensemble_size, -1, 3)
if step_scaler is None:
    ...  # batched fast path, early return
return self._run_guided_step_per_particle(coords_4d, model, sampler, features, context, step_scaler, num_particles)

Pure extraction — no behavior change. The guided body (per-particle sampler.step, loss/log-correction reduction, torch.stack restacking) is moved verbatim.

Why the two paths stay separate

The issue notes the guided/unguided paths look duplicated, but they can't be fully merged: the unguided path runs a single batched forward pass over all particles * ensemble structures (one sampler.step), while the guided path must step each particle independently so its guidance gradient is computed only from its own state (correct FK semantics). Funneling the batched output through the per-particle stacker would misinterpret the whole batch as one particle. So the split is intentional; this change makes that boundary explicit and readable rather than collapsing it. Documented the rationale in the new method's docstring.

Closes #77.

Summary by CodeRabbit

• Refactor
  • Improved the stepping flow for guided sampling by separating the per-particle path into a dedicated internal routine.
  • Preserved the existing fast path when no step scaling is used, including returning denoised coordinates when available.
  • Kept guided particle-by-particle behavior unchanged, with results still aggregated and reshaped into the expected output format.

[coderabbitai]

Note

Currently processing new changes in this PR. This may take a few minutes, please wait...

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 735346ef-f564-49e2-90fc-116e10ac8364

📥 Commits

Reviewing files that changed from the base of the PR and between 38f244a and 8d4be87.

📒 Files selected for processing (1)

• src/sampleworks/core/scalers/fk_steering.py

 __________________________________________________________________________________
< This is the beginning of a beautiful friendship...between you and bug-free code. >
 ----------------------------------------------------------------------------------
  \
   \   \
        \ /\
        ( )
      .( o ).

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[Copilot]

Pull request overview

This PR refactors FKSteering._run_step to separate the guided per-particle stepping logic into a dedicated helper, improving readability while preserving the existing split between the unguided batched fast path and the guided per-particle FK semantics.

Changes:

• Extracted the guided per-particle loop from _run_step into _run_guided_step_per_particle.
• Simplified _run_step to a small dispatch between the unguided batched fast path and guided per-particle path.
• Added a NumPy-style docstring documenting why the two paths cannot be fully merged (per-particle gradient semantics).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.