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

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

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

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

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

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Co-authored-by: xraymemory <me.anzuoni@gmail.com>
Co-authored-by: M E A <xraymemory@users.noreply.github.com>

Author: 2 people

src/sampleworks/core/scalers/fk_steering.py

@@ -286,6 +286,52 @@ def _run_step(
 
             return step_output, denoised_4d
 
+        return self._run_guided_step_per_particle(
+            coords_4d, model, sampler, features, context, step_scaler, num_particles
+        )
+
+    def _run_guided_step_per_particle(
+        self,
+        coords_4d: torch.Tensor,
+        model: FlowModelWrapper,
+        sampler: TrajectorySampler,
+        features: GenerativeModelInput,
+        context: StepParams,
+        step_scaler: StepScalerProtocol,
+        num_particles: int,
+    ) -> tuple[SamplerStepOutput, torch.Tensor | None]:
+        r"""Run the denoising step one particle at a time, then restack to a batch.
+
+        Each particle is stepped independently so its guidance gradient is computed
+        only from its own state (correct FK semantics). This is the guided
+        counterpart to the batched fast path in :meth:`_run_step`; the two are kept
+        separate because the fast path runs a single forward pass over all particles,
+        which cannot provide per-particle gradients.
+
+        Parameters
+        ----------
+        coords_4d : torch.Tensor
+            Current coordinates, shape ``(particles, ensemble, atoms, 3)``.
+        model : FlowModelWrapper
+            Model wrapper for denoising.
+        sampler : TrajectorySampler
+            Sampler that handles the step mechanics.
+        features : GenerativeModelInput
+            Model features/inputs.
+        context : StepParams
+            Step context with time info and optionally reward.
+        step_scaler : StepScalerProtocol
+            Active step scaler providing per-particle guidance.
+        num_particles : int
+            Number of particles to iterate over.
+
+        Returns
+        -------
+        tuple[SamplerStepOutput, torch.Tensor | None]
+            ``(step_output, denoised_4d)`` matching :meth:`_run_step`, with
+            ``step_output.state`` shape ``(particles * ensemble, atoms, 3)`` and
+            ``denoised_4d`` shape ``(particles, ensemble, atoms, 3)``.
+        """
         states_per_particle: list[torch.Tensor] = []
         denoised_per_particle: list[torch.Tensor] = []
         losses_per_particle: list[torch.Tensor] = []