Param as dict

[pmelchior]

This PR brings significant changes under the hood in who parameters are stored and treated. In particular:

• Parameters now derive from dict, so allow lookup like parameters["spectrum_0"].
• The global namespace hold a parameter_registry, which uses the memory location of a module instance to hold the parameters of that instance. This entry is created by the init function Parameters(some_module_instance). The parameter set can be retrieved by calling the property Module.parameters. Module also allows for set and get of the parameter values.
• The new free functions infer.fit and infer.sample now accept parameters from multiple modules and internally combine them into a dynamically created eqx.Module that contains attributes for every key/value pair in the parameters dict. Only this new module of values get optimized/sampled, which massively reduces the overhead per iteration as it avoids copying non-optimized data structures.

The main purpose is to allow Parameters to be defined not just for Scene, but for any other Module, in particular Observation. As an example, we can adjust the astrometric offset by parameterizing the .shift attribute in the ConvolutionRenderer:

with sc2.Parameters(obs_hsc):
    sc2.Parameter(obs_hsc.renderer[1].shift, name="shift", stepsize=1e-3 * u.arcsec)

A detailed example is given in the multi-observation howto.

As a consequence of these changes, the signature of the fit/sample functions does not require the parameters argument (it's held already by the module that gets optimized). This shorter call signature will do and is now used in the docs:

with sc2.Parameters(scene): # note: no "as parameters"
    ...
scene.fit(obs, max_iter=1000) # note no parameters argument

While the changes are quite substantial, scripts that are based on the v0.3 API should all still work.

[github-actions]

Before [b2c091d]	After [49faf3b]	Ratio	Benchmark (Parameter)
0	0	n/a	bbox_benchmarks.BoxSuite.mem_bbox_contains(16)
0	0	n/a	bbox_benchmarks.BoxSuite.mem_bbox_contains(2)
0	0	n/a	bbox_benchmarks.BoxSuite.mem_bbox_contains(2048)
0	0	n/a	bbox_benchmarks.BoxSuite.mem_bbox_contains(256)
0	0	n/a	bbox_benchmarks.mem_bbox_creation
35.7±0.4μs	36.4±0.8μs	1.02	bbox_benchmarks.BoxSuite.time_bbox_contains(2048)
7.75±0.1μs	7.88±0.1μs	1.02	bbox_benchmarks.time_bbox_creation
35.2±0.3μs	35.5±0.4μs	1.01	bbox_benchmarks.BoxSuite.time_bbox_contains(16)
33.6±1μs	34.1±0.2μs	1.01	bbox_benchmarks.BoxSuite.time_bbox_contains(2)
35.9±0.4μs	35.6±0.3μs	0.99	bbox_benchmarks.BoxSuite.time_bbox_contains(256)

Click here to view all benchmarks.