Skip to content

JointFitter and two examples. - #288

Open
KriSun95 wants to merge 13 commits into
sunpy:mainfrom
KriSun95:joint-fitting
Open

JointFitter and two examples.#288
KriSun95 wants to merge 13 commits into
sunpy:mainfrom
KriSun95:joint-fitting

Conversation

@KriSun95

@KriSun95 KriSun95 commented Jul 22, 2026

Copy link
Copy Markdown
Collaborator

Adds the first version of the new JointFitter and two examples for the example gallery.

This is still a work in progress. My local docs building is not working.

To add:

  • changelog
  • get local docs build working
  • make tests
  • run pre-commit

Thoughts

This is workable but there are still somethings to do or clean-up. I'm not sure if the following should be this PR or another one.

  • We have to decide how to choose and optimiser and statistic. At the minute, JointFitter needs to be initialised with one of each since I'm trying to follow their set-up. The issue is, not all of their other fitters follow this and end up just bypassing this and hardcoding in their optimiser and statistic choices anyway. How do we want to do this?
  • Unit support! Right now, JointFitter just handles units briefly and crudely to get stuff to work. What I've written isn't complicated and we should look into altering the Astropy decorator that handles units.
  • The thick target function has its issues with input+output units and also returning NaNs for some very high precision input values. I've not included the thick target in any example here for that reason. This should be another PR but I wasn't sure if that should be one on its own or with one of Jake's PRs. I just wanted to mention it here.

@jajmitchell , @samaloney , let me know what you both think about this. I think the JointFitter doc-strings could use a little work but other than that, what would be crucial to get this initial stuff merged in your eyes?

@settwi

settwi commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

Nice. I want to add a sensible “what if” to think about.

What if a user wants to use a prior distribution that is not a uniform prior? Bounding the parameters is equivalent to a uniform prior. Say I want to use a truncated normal distribution on a parameter, or a Jeffreys prior. How hard would that be to implement as this code is written, or as Astropy modeling is implemented?

@KriSun95

Copy link
Copy Markdown
Collaborator Author

@settwi , that's a good point and worth thinking about.

The permanent fix would be to PR Astropy

Personally, I think that the place to allow this would be in the module level fitter_to_model_params_array function but in Astropy main. This function is the one responsible for building up the parameter list for a model from the partial list of fitted parameters and is the one that enforces the bounds in the first place.

If we could just allow the bounds attribute for the parameter object to accept a prior function of the user's choice (either in place of a simple length 2 tuple or as well as) then I think the code-block that enforces the given bounds at the minute (here) could just run the value through the parameter's "bound function" to get the corresponding value.

This would make the bounds on any parameter far more powerful for any fitting.

The temporary fix (while we wait for Astropy to accept)

I can think of two ways:

  1. We can edit our fitter_to_model_params_array to do this.
  2. Add a prior attribute to our model parameters (like temperature, etc.) that we can check for once the parameters are returned from fitter_to_model_params_array in our own fitter objects.

@samaloney samaloney left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is really nice!

I wonder would it make sense to extract an ABC which defines the API then make this a concrete implementation?

The tests cover:

  • unitless models,
  • forward-direction tying
    but misses:
  • fixed params
  • bounds enforcement check,
  • weights,
  • Quantity units.

The units-handling branch and the weights/fkwarg paths used in are untested. I do think tests for any feature expose by the API would be needed - could be done in a later PR but would be better here.

It would be good to try and find some specific examples of the thick target function issues so we can pin it down and hopefully fix.

On the optimiser and statistic I think the way it set up now makes sense if need to define our own API and then map the the underling fitters - other subclasses can override the methods and raise error is not supported

Comment on lines +205 to +215
## should really call self._opt_method()
# optimize.least_squares just minimises the square of what comes out of self.objective_function
# need the lambda function for now since `minimize` won't accept other kwargs
fun = lambda x: self.objective_function(x, *(models, farg), **fkwarg) # noqa
result = minimize(
fun,
jmodel_params,
bounds=tuple(zip(jmodel_bounds[0], jmodel_bounds[1])),
method="Nelder-Mead",
tol=1e-8,
)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As partially noted in the PR description the __init__ validates and stores self._opt_method/self._stat_method here always calls scipy.optimize.minimize(..., method="Nelder-Mead") with a hardcoded sum-of-squares.

Both gallery examples call JointFitter(optimizer=SLSQP, statistic=leastsquare) which doesn't actually use the optim/stats functions.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Absolutely, and I don't like this either. I wanted to get the code up here to discuss first but I definitely thought this is something to fix. I'm hoping my pseudo code gets at this and I'm happy to build on this.


for mod_num, model in enumerate(models):
# units, make sure we have units
if parameter_units is not None:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

mod_params is only assigned when parameter_units is not None. It's unreachable via the public call path (_get_param_units never returns None for non-empty model lists)



## KRIS: All I did was add ``model_list``
def fitter_to_model_params_array(model, fps, use_min_max_bounds=True, *, fit_param_indices=None, model_list=None):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should see if we can eventually PR this to astropy

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Absolutely. I Think that pretty much all the JointFitter code and fitters.py code as it stands could live in Astropy eventually. The only thing I've essentially done (after a lot of work to get there) was to:

  • Get the fitter to accept (model, x, y) groups instead of just one
    • This mainly meant I had to add in some loops to loop through the groups (+minimal other things)
  • Make those two module level functions (only because one of these functions uses the other) accept this model_list parameter

I don't see why the other fitting classes in Astropy couldn't just be made to do this. Then we only have to have our own fitters like the NelderMeadLeastSquaresFitter fitter with our own optimiser and statistics modules to let us do whatever we want.

@KriSun95

Copy link
Copy Markdown
Collaborator Author

This is really nice!

I wonder would it make sense to extract an ABC which defines the API then make this a concrete implementation?

Thanks! I've been happy with how it's turned out so far.

I would love to do that! It's something that the API in Astropy isn't very solid in and I don't like exactly how I've done it here.

The way Astropy does it is by defining the base Fitter class which any other fitter needs to inherit from (I also inherit from Fitter for our JointFitter). These other classes then redefine the same functions as I did, like objective_function, _run_fitter, and __call__. The thing that annoys me with the Astropy fitters is that a lot of them inherit from Fitter but never actually use it (i.e., they never call super in their init) then just hard code everything.

I can try and make JointFitter a sort of base class that needs to be initialised with an optimiser and statistic. Then inherit from it in a new class, like NelderMeadLeastSquaresFitter, that we can use in the same way. This might mean we can just be a bit more rigid the Astropy in how new fitters should be created.

E.g.,
In fitters.py:

from scipy.optimize import minimize
from some_module import least_squares_func

class JointFitter(Fitter):
    def __init__(self, optimizer, statistic):
        super().__init__(optimizer=optimizer, statistic=statistic)

    def joint_model_to_fit_params(self, models):
        ...

    def _update_model_params(self, models, fps, jfit_param_indices):
        ...

    def _verify_input(self, args):
        ...

    @staticmethod
    def _get_param_units(models):
        ...

    def _run_fitter(self, models, farg, fkwarg=None):
        raise NotImplementedError("Subclasses should implement this method.")

    def objective_function(self, models, farg, fkwarg=None):
        raise NotImplementedError("Subclasses should implement this method.")

    def __call__(self, *args, fkwarg=None, inplace=False):
        # will check inputs then return output of `self._run_fitter`
        ...
        return self._run_fitter(...)


class NelderMeadLeastSquaresFitter(JointFitter):
    def __init__(self):
        super().__init__(optimizer=minimize, statistic=least_squares_func)

    def objective_function(self, fps, *args, weights=None, jfit_param_indices=None, parameter_units=None):
        # gathers parameters and evaluates model to run`self._stat_method` which is `least_squares_func` here
        ...
        self._stat_method()

    def _run_fitter(self, models, farg, fkwarg=None):
        # sets up and runs `self._opt_method` which is `minimize` here
        ...
        self._opt_method(...)
        ...

The user could then just:

fitter = NelderMeadLeastSquaresFitter()
result = fitter(model0, x0, y0, ..., modelN, xN, yN)

The JointFitter class could realistically replace the base Fitter class from Astropy eventually...I think.

The tests cover:

  • unitless models,
  • forward-direction tying
    but misses:
  • fixed params
  • bounds enforcement check,
  • weights,
  • Quantity units.

The units-handling branch and the weights/fkwarg paths used in are untested. I do think tests for any feature expose by the API would be needed - could be done in a later PR but would be better here.

Agreed! I'll work through this list and update the tests to account for them.

It would be good to try and find some specific examples of the thick target function issues so we can pin it down and hopefully fix.

Agreed again and I'm working on another branch to do this.

On the optimiser and statistic I think the way it set up now makes sense if need to define our own API and then map the the underling fitters - other subclasses can override the methods and raise error is not supported

Have I answered this with the pseudo code above?

@KriSun95

KriSun95 commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator Author

New tests to add:

  • unitless models,
  • forward-direction tying
  • fixed params
  • bounds enforcement check,
  • weights,
  • Quantity units

@KriSun95

Copy link
Copy Markdown
Collaborator Author

Don't believe the failed ubuntu-latest test. I've seen it fail then pass with no change to the code here and in PR #289

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants