automatic wrappers for grid search and grid search with optimization

[benjaminpope]

Many of our functions are a grid search, or a grid search followed by optimization.

I would like to abstract these with unctions that you pass a loss function and sample_dict and it will do the grid search and optimization appropriately.

One point to consider is how to handle the names of parameters that we are BFGS-ing over (ie flux), and which are coordinates to be fixed (dra, ddec) - these functions should be fully general. The implementation is easy either way but what API do we pick - perhaps loss_param = 'flux' or some syntax like that to pick out which one we are optimizing?

[LouisDesdoigts]

Not sure I fully follow the problem here. If we aren't grid-searching/optimising some parameters, they wouldn't be in the sample_dict right? Everything else that is static should just be taken from as the value within the object no?

This would be the same way numpyro works, you don't have to specify what is fixed you only specify what is sampled over, everything else is treated as static and its value is used by the model object implicitly.

[benjaminpope]

An automatic grid search with optimization is now implemented, for cartesian grids.

I want to suggest that a model object has an attribute which lists which parameters are considered as coordinates to be mapped over in grid searches, because they may have many other parameters to be inferred but which are not in general going to require grid searches (and which would scale badly at high dimension).

I suggest a syntax we just do

self.coords = ['dra', 'ddec'] # ['sep', 'pa'], or disk spatial params, etc

The other parameters, frozen at grid search time, would be in sample_dict as scalars. Is there anything wrong with doing this?

I also suggest that we have a helper attribute self.contrast whose getter returns 1/self.flux (which is a mandatory keyword for all sources), but whose setter actually updates self.flux = 1/contrast so that we can interface with this nicely.

[LouisDesdoigts]

I'm not really convinced of this, or maybe I don't really understand where things are happening.

Where is this inference of the other parameters happening? Are these other parameters being inferred inside this grid search function? If not I can't see why they should live in the sample_dict as they are already inside the class? Bit hard to say much more on this without actually knowing what these functions are doing and how they fit into the wider workflow.

I am in general not a fan of having leaves within a class that are used to refer to other leaves within the class, I've not really found a situation where that is warranted, but again, I don't have the context to be able to provide any concrete examples of where this might be troublesome. You could hold the sample_dict as part of the model itself, but I would tend to lean towards creating a Grid or ParametersGrid class, that stores the sample dict and grid_likelihood or whatever, than can be a leaf of an arbitrary model. This lets you more rigidly define the parameters that are inferred (within the model class) and those that are grid searched (in the Grid class), providing a concrete boundary between them and without any self-referential stuff. This would be a more more equinox'y approach IMO and gives you more flexibility in terms of how they interact with each other.

---

Parametric getter method for contrast 👍

Parametric setter method for contrast 👎

The getter methods are totally fine and just make working with classes easier. However, if you create a setter method that completely breaks with how everything else is done. That one (faux) leaf requires its whole own setter method, that you then won't be able to update using zodiax methods. This is a fundamental limit because equinox updates leaves based on a where function, and you obviously can't point it to something that doesn't exist. You can use a __getattr__ method to point it to flux, but then you have no way to apply the inverse, which again leave you at either creating a specific setter function for that leaf alone, and just breaking from the whole ability to use zodiax.

I think you could mayyyybe get this behavior by taking a similar approach to eqx.nn.Shared layer and having custom behavior in the __setattr__ method of the shared class equivalent, but I can neither confirm nor deny this, and am not 100% on how this would interact with zodiax methods.

[LouisDesdoigts]

Oh and looking at #20, I think the potential parametric setter solution would also break methods like __mul__, __add__ etc, unless you hard-code specific behavior. You can kind of see how this gets out of hand from a code and maintenance perspective when you can simply just do model.set('flux', 1/contrast)