Optimizer and Scheduler Type Hinting

[cedriclim1]

What problem does this PR address?

Optimizer and scheduler parameter initialization has been using dictionaries making type-hinting difficult. This PR implements a nested class in core/ml/optimizer_mixin.py with dataclasses to allow for clean type-hinting.

What should the reviewer(s) do?

• Check implementations of various optimizer and scheduler types.
• Check notebook that everything works.
• (Optional) quick implementation in Ptychography to make sure there are no issues. Note: I had to make sure my tomography_opt.py file would allow for OptimizerType changes.

Example to the last checkbox:

    @optimizer_params.setter
    def optimizer_params(self, d: dict[str, OptimizerType]):
        """Set the optimizer parameters."""
        if isinstance(d, (tuple, list)):
            d = {k: {} for k in d}

        targets = {
            "object": self.obj_model,
            "pose": self.dset,
        }

        for k, v in d.items():
            if k not in targets:
                raise ValueError(f"Unknown optimization key: {k}")

            if not isinstance(v, OptimizerType):
                if "lr" not in v:
                    v["lr"] = self._get_default_lr(k)

            targets[k].optimizer_params = v

Notebook

opt_sched_type_hinting.ipynb

[arthurmccray]

Overall looks good, I like this solution :) I just promoted the Params in the namespace and added docstrings. The only question/thing left is to do with SchedulerParams. Do you think we should have the params() method take the num_steps in addition to base_LR? There are two reasons for doing this:

1. allow CosineAnnealing.T_max to be int | None = None, so you don't have to specify the number of steps in two places. It would just have the same if self.T_max is None logic like is done for base_LR
2. This way the factor attribut of Exponential would work. ExponentialLR only takes a gamma arg, and as written in breaks if trying to pass factor. This way, factor would be used to set gamma based on the num_steps when calling params().
   (also it could be num_iters or whatever, i don't really care what the arg name is)

I will also pull into diffractive_imaging to make any changes required for ptycho (but cuz we're allowing backwards compatibility it shouldn't break anything? might just make the type checker angry)

[cedriclim1]

@arthurmccray re: #186 (comment)

Added back in num_iters as an input to .params. So the convention I chose here, if T_max or num_iter are already instantiated in the class, it will always use that value overriding whatever has been input into .params. For example:

sched = CosineAnnealing(T_max = 100)

schedulers = self._set_scheduler_params.params(LR, 10)

sched.T_max = 100

Happy to change this convention though.

[arthurmccray]

lgtm

[arthurmccray]

Tested with iterative ptycho and some notebooks and seems to be working, thanks for adding the tests.