Optimizer changes

[cedriclim1]

Summary of changes

Bug fixes

1. Shared probe origin drifted in the serial fit even with disk/lattice LRs zeroed

OriginND was registered as a submodule of DiskTemplate, SyntheticDiskLattice, and GaussianBackground. Each component's optimizer included origin.coords, so the serial fit stepped it three times per iteration. Setting disk0.lr=0 and lat0.lr=0 zeroed
two updates, but gaussianbackground.lr=1e-5 kept moving the probe.

2. Pre-existing upstream 'Adam' object is not iterable in _rebuild_optimizer_after_trainability_change

dict(self.optimizer_params) and dict(self.scheduler_params) in the infer*_rebuild_params helpers tried to iterate over OptimizerParams.Adam / SchedulerParams.Plateau dataclasses, which aren't iterable. This blocked set_component_trainable(...).

Files changed

quantem/src/quantem/core/fitting/base.py

• OriginND now inherits OptimizerMixin and owns the single optimizer for coords (with DEFAULT_OPTIMIZER="adamw", DEFAULT_LR=1e-2, initialize_optimizer(...) mirroring RenderComponent).
• RenderComponent.get_optimization_parameters excludes any params belonging to a shared OriginND submodule.
• AdditiveRenderModel gained _origin_fallback_lr and _initialize_origin_optimizer, called from initilize_independant_optimizers. The three loop helpers — zero_grad_optimizers, step_optimizers, step_schedulers — also drive the origin optimizer.
  set_independant_optimizer_params and rebuild_independant_optimizers handle "origin" too.
• _infer_optimizer_rebuild_params / _infer_scheduler_rebuild_params now use a dataclass-aware conversion ({f.name: getattr(op, f.name) for f in dataclasses.fields(op)} with _name → "type").
• fit_render docstring documents the new "origin" key in optimizer_params / scheduler_params.

quantem/src/quantem/diffraction/model_fitting.py

• _BatchedPlan.from_model honors optimizer_params["origin"] first, then falls back to lat → gaussbg → 1e-2 (matches the serial-path fallback). Adds "origin" to component_keys so scheduler-by-name and frozen_components=["origin"] dispatch reach it.

Notebooks/20260511-multi-diff-batched.ipynb

• Cell 5e1689dd: added "origin": {"type": "adamw", "lr": 0} to individual_optimizer_params to lock the probe explicitly.

Scripts/smoke_origin_freeze.py (new)

• 4 checks: serial freeze (origin.lr=0 → bit-exact), serial moves (origin.lr=1e-3 → diff > 1e-6), batched freeze, legacy fallback (no "origin" key but lat0.lr=0). All pass.
  Public API

optimizer_params={"origin": {"type": "adamw", "lr": 0.0}} (and scheduler_params={"origin": {...}}) is the new way to control the probe origin across all components in both the serial and batched paths. When "origin" is omitted, LR falls back to the
first origin-owning component in priority order SyntheticDiskLattice > GaussianBackground > DiskTemplate, which preserves behavior for notebooks that previously relied on lat0.lr implicitly driving the probe.