Astropy - Modeling

In [1]:
import numpy as np
%matplotlib inline
import matplotlib.pyplot as plt
In [2]:
from astropy.modeling import models, fitting

Setting up models

In [3]:
g = models.Gaussian1D(amplitude=1.2, mean=0.9, stddev=0.5)
print(g)

Model: Gaussian1D
Inputs: ('x',)
Outputs: ('y',)
Model set size: 1
Parameters:
    amplitude mean stddev
    --------- ---- ------
          1.2  0.9    0.5

Model parameters can be accessed as attributes:

In [4]:
g.amplitude
Out[4]:
Parameter('amplitude', value=1.2)
In [5]:
g.mean
Out[5]:
Parameter('mean', value=0.9)
In [6]:
g.stddev
Out[6]:
Parameter('stddev', value=0.5)

and can also be updated via those attributes:

In [7]:
g.amplitude = 0.8
g.amplitude
Out[7]:
Parameter('amplitude', value=0.8)

Models can be evaluated directly (so they are useful without even fitting):

In [8]:
g(1.3)
Out[8]:
0.5809192296589527
In [9]:
x = np.linspace(-5, 5, 1000)
plt.plot(x, g(x))
Out[9]:
[<matplotlib.lines.Line2D at 0x109fadf60>]

Fitting models to data

We start off by generating fake data:

In [10]:
np.random.seed(0)
x = np.linspace(-5., 5., 200)
y = 3 * np.exp(-0.5 * (x - 1.3)**2 / 0.8**2)
y += np.random.normal(0., 0.2, x.shape)
plt.plot(x, y, 'ko')
Out[10]:
[<matplotlib.lines.Line2D at 0x10a008d68>]

We set up a fitter:

In [11]:
fitter = fitting.LevMarLSQFitter()

We now fit the data using a trapezoid model and a Gaussian

In [12]:
t_init = models.Trapezoid1D(amplitude=1., x_0=0., width=1., slope=0.5)
t = fitter(t_init, x, y)

g_init = models.Gaussian1D(amplitude=1., mean=0, stddev=1.)
g = fitter(g_init, x, y)

Finally we can plot the data and the two models

In [13]:
plt.figure(figsize=(8,5))
plt.plot(x, y, 'ko')
plt.plot(x, t(x), 'b-', lw=2, label='Trapezoid')
plt.plot(x, g(x), 'r-', lw=2, label='Gaussian')
plt.xlabel('Position')
plt.ylabel('Flux')
plt.legend(loc=2)
Out[13]:
<matplotlib.legend.Legend at 0x10b09bcf8>

Multi-dimensional models

The modeling framework is not restricted to 1-d models and datasets. To demonstrate this, we create a 2-d dataset:

In [14]:
np.random.seed(0)
y, x = np.mgrid[:128, :128]
z = 2. * x ** 2 - 0.5 * x ** 2 + 1.5 * x * y - 1.
z += np.random.normal(0., 0.1, z.shape) * 50000.
plt.imshow(z, origin='lower', interpolation='nearest', vmin=-1e4, vmax=5e4)
Out[14]:
<matplotlib.image.AxesImage at 0x10b409080>

As before, we set up a fitter:

In [15]:
fitter = fitting.LinearLSQFitter()

We create the initial model:

In [16]:
p_init = models.Polynomial2D(degree=2)

and we can now fit the model to the data:

In [17]:
p = fitter(p_init, x, y, z)

We can now look at the best model

In [18]:
plt.imshow(p(x, y), origin='lower', interpolation='nearest', vmin=-1e4, vmax=5e4)
Out[18]:
<matplotlib.image.AxesImage at 0x10bc80710>

and the residual

In [19]:
plt.imshow(z - p(x, y), origin='lower', interpolation='nearest', vmin=-1e4, vmax=5e4)
Out[19]:
<matplotlib.image.AxesImage at 0x10bec17b8>

Compound models

In some cases, we may want to fit multiple models at the same time (for example two Gaussians, or a background polynomial and a Gaussian, etc.). This is now possible in Astropy 1.0 and later. As before we start off by generating synthetic data:

In [20]:
np.random.seed(42)
g1 = models.Gaussian1D(1, 0, 0.2)
g2 = models.Gaussian1D(2.5, 0.5, 0.1)
x = np.linspace(-1, 1, 200)
y = g1(x) + g2(x) + np.random.normal(0., 0.2, x.shape)
plt.plot(x, y, 'ko')
Out[20]:
[<matplotlib.lines.Line2D at 0x10b1df390>]

Now we fit this with a combination of two Gaussians:

In [21]:
gg_init = models.Gaussian1D(1, 0, 0.1) + models.Gaussian1D(2, 0.5, 0.1)
fitter = fitting.SLSQPLSQFitter()
gg_fit = fitter(gg_init, x, y)

Optimization terminated successfully.    (Exit mode 0)
            Current function value: 6.83285936044
            Iterations: 14
            Function evaluations: 137
            Gradient evaluations: 14

Finally, we can plot the best-fitting model:

In [22]:
# Plot the data with the best-fit model
plt.figure(figsize=(8,5))
plt.plot(x, y, 'ko')
plt.plot(x, gg_fit(x), 'r-', lw=2)
plt.xlabel('Position')
plt.ylabel('Flux')
Out[22]:
<matplotlib.text.Text at 0x10c1266a0>

Exercises

Level 1

Read in the data from data/fitting_data.txt and try and fit it with a straight line and a parabola, then make a plot of the resulting fits.

Level 2

Download the first FITS spectrum from this page, read it in with astropy.io.fits and try and fit a compound model of a gaussian plus a constant to the spectrum (just assume the points are equally spaced in wavelength space, even if that may not be true). Hint: you may have to play around with setting better initial values for the position of the line.

Level 3

Try and read in an image of a star field, and try fitting a 2D Gaussian to one of the stars.