Implement parallel fitting

[astrofrog]

This currently implements a single function that can be used to fit models to N-dimensional cubes. At the moment the internal code is very much a prototype and doesn't implement a number of features. But it should in principle already be possible to pass in an N-dimensional cube and fit an M-dimensional model to a subset of the axes, for example fit a 1D model to spectra in a spectral cube, or 2D Gaussians to the celestial axes and so on.

Simple example

    import numpy as np
    from astropy.modeling.models import Gaussian1D
    from astropy.modeling.fitting import LMLSQFitter
    from astropy.modeling.fitting_parallel import parallel_fit_model_nd

    NWAV = 100
    NMOD = 200

    wav = np.linspace(10, 20, NWAV)
    amplitude = np.linspace(1, 2, NMOD)
    mean = np.linspace(12, 18, NMOD)
    sigma = np.linspace(1, 2, NMOD)

    data = amplitude * np.exp(-(wav[:, None] - mean) ** 2 / 2 / sigma **2)
    data += np.random.random(data.shape) / 5

    g = Gaussian1D(amplitude=2, mean=np.ones(NMOD), stddev=np.ones(NMOD))

    from tqdm.dask import TqdmCallback
    with TqdmCallback(desc="fitting"):

        g_fit = parallel_fit_model_nd(
            model=g,
            fitter=LMLSQFitter(),
            data=data,
            fitting_axes=0,
            world={0: wav}
        )

    import matplotlib.pyplot as plt

    fig = plt.figure()
    ax = fig.add_subplot(3, 1, 1)
    ax.text(0.05, 0.8, 'amplitude', transform=ax.transAxes)
    ax.plot(amplitude)
    ax.plot(g_fit.amplitude.value)
    ax = fig.add_subplot(3, 1, 2)
    ax.text(0.05, 0.8, 'mean', transform=ax.transAxes)
    ax.plot(mean)
    ax.plot(g_fit.mean.value)
    ax = fig.add_subplot(3, 1, 3)
    ax.text(0.05, 0.8, 'stddev', transform=ax.transAxes)
    ax.plot(sigma)
    ax.plot(g_fit.stddev.value)
    fig.savefig('results.png')

Things that still need to be done

• [x] Support passing in a WCS for the world coordinates of the fitting axes
• [x] Support passing in N-dimensional arrays for the world coordinates of the fitting axes
• [x] Support not passing in world and just using pixel coordinates for axes
• [x] Add some form of error handling if the fit fails
• [x] Add the ability to optionally visualize failing or all fits
• [x] Add extensive tests
• [x] Test out with dask.distributed
• [x] Check if works ok with multi-threading even if not much faster (i.e. make sure it is thread-safe at least)
• [x] Passing parameters to fitters
• [x] Try and re-write dask part to use single array with original dimensionality/chunks instead of all the raveling etc.
• [x] Try out with real use cases
• [x] Optimize performance as-needed
• [x] Support passing in weights if fitter supports it
• [x] Generalize diagnostics to take a callback function rather than do plotting
• [x] Add tests for handling of units
• [ ] Write user documentation

Related issues/PRs:

• https://github.com/astropy/astropy/issues/16593
• https://github.com/dask/dask/issues/11188

Open questions:

• [x] When passing in a WCS, should we attach units to the world coordinates and then expect that the model will have units attached to parameters?
• [x] How do we handle units on the data, should we pass in quantity? What about when data is a dask array?
• [x] What should default scheduler be?
• [x] What do we actually want to do for diagnostics? Do we want to provide a way to have a callback function instead of hard-coding some plotting code?

Beyond this PR:

• Could support, either as an option or default, the ability to return an uncomputed dask array (what would be returned though, since normally it should be a model? Would it be a model with dask array parameters? or something else?)
• Add ability to specify a mask of which indices to fit (for example if fitting spaxels in a spectral cube, which spaxels to include/not include)

[github-actions[bot]]

Thank you for your contribution to Astropy! 🌌 This checklist is meant to remind the package maintainers who will review this pull request of some common things to look for.

• [ ] Do the proposed changes actually accomplish desired goals?
• [ ] Do the proposed changes follow the Astropy coding guidelines?
• [ ] Are tests added/updated as required? If so, do they follow the Astropy testing guidelines?
• [ ] Are docs added/updated as required? If so, do they follow the Astropy documentation guidelines?
• [ ] Is rebase and/or squash necessary? If so, please provide the author with appropriate instructions. Also see instructions for rebase and squash.
• [ ] Did the CI pass? If no, are the failures related? If you need to run daily and weekly cron jobs as part of the PR, please apply the "Extra CI" label. Codestyle issues can be fixed by the bot.
• [ ] Is a change log needed? If yes, did the change log check pass? If no, add the "no-changelog-entry-needed" label. If this is a manual backport, use the "skip-changelog-checks" label unless special changelog handling is necessary.
• [ ] Is this a big PR that makes a "What's new?" entry worthwhile and if so, is (1) a "what's new" entry included in this PR and (2) the "whatsnew-needed" label applied?
• [ ] At the time of adding the milestone, if the milestone set requires a backport to release branch(es), apply the appropriate "backport-X.Y.x" label(s) before merge.

[github-actions[bot]]

👋 Thank you for your draft pull request! Do you know that you can use [ci skip] or [skip ci] in your commit messages to skip running continuous integration tests until you are ready?

[astrofrog]

Might also want to have a way to specify a mask to say which pixels to fit vs ignore

[astrofrog]

Closing as the astropy PR is now open.