File: test_stepmodel.py

package info (click to toggle)
lmfit-py 1.3.3-4
  • links: PTS, VCS
  • area: main
  • in suites: sid, trixie
  • size: 2,332 kB
  • sloc: python: 13,071; makefile: 130; sh: 30
file content (105 lines) | stat: -rw-r--r-- 2,987 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
import numpy as np

from lmfit.models import ConstantModel, RectangleModel, StepModel


def get_data():
    np.random.seed(2021)
    x = np.linspace(0, 10, 201)
    dat = np.ones_like(x)
    dat[:48] = 0.0
    dat[48:77] = np.arange(77-48)/(77.0-48)
    dat = dat + 5e-2*np.random.randn(len(x))
    dat = 110.2 * dat + 12.0
    return x, dat


def test_stepmodel_linear():
    x, y = get_data()
    stepmod = StepModel(form='linear')
    const = ConstantModel()
    pars = stepmod.guess(y, x)
    pars = pars + const.make_params(c=3*y.min())
    mod = stepmod + const

    out = mod.fit(y, pars, x=x)

    assert out.nfev > 5
    assert out.nvarys == 4
    assert out.chisqr > 1
    assert out.params['c'].value > 3
    assert out.params['center'].value > 1
    assert out.params['center'].value < 4
    assert out.params['sigma'].value > 0.5
    assert out.params['sigma'].value < 3.5
    assert out.params['amplitude'].value > 50


def test_stepmodel_erf():
    x, y = get_data()
    stepmod = StepModel(form='linear')
    const = ConstantModel()
    pars = stepmod.guess(y, x)
    pars = pars + const.make_params(c=3*y.min())
    mod = stepmod + const

    out = mod.fit(y, pars, x=x)

    assert out.nfev > 5
    assert out.nvarys == 4
    assert out.chisqr > 1
    assert out.params['c'].value > 3
    assert out.params['center'].value > 1
    assert out.params['center'].value < 4
    assert out.params['amplitude'].value > 50
    assert out.params['sigma'].value > 0.2
    assert out.params['sigma'].value < 1.5


def test_stepmodel_stepdown():
    x = np.linspace(0, 50, 201)
    y = np.ones_like(x)
    y[129:] = 0.0
    y[109:129] = 1.0 - np.arange(20)/20
    y = y + 5e-2*np.random.randn(len(x))
    stepmod = StepModel(form='linear')
    pars = stepmod.guess(y, x)

    out = stepmod.fit(y, pars, x=x)

    assert out.nfev > 10
    assert out.nvarys == 3
    assert out.chisqr > 0.2
    assert out.chisqr < 5.0
    assert out.params['center'].value > 28
    assert out.params['center'].value < 32
    assert out.params['amplitude'].value > 0.5
    assert out.params['sigma'].value < -2.0
    assert out.params['sigma'].value > -8.0


def test_rectangle():
    x = np.linspace(0, 50, 201)
    y = np.ones_like(x) * 2.5
    y[:33] = 0.0
    y[162:] = 0.0
    y[33:50] = 2.5*np.arange(50-33)/(50-33)
    y[155:162] = 2.5 * (1 - np.arange(162-155)/(162-155))
    y = y + 5e-2*np.random.randn(len(x))
    stepmod = RectangleModel(form='linear')
    pars = stepmod.guess(y, x)

    out = stepmod.fit(y, pars, x=x)

    assert out.nfev > 10
    assert out.nvarys == 5
    assert out.chisqr > 0.2
    assert out.chisqr < 5.0
    assert out.params['center1'].value > 8
    assert out.params['center1'].value < 14
    assert out.params['amplitude'].value > 2.0
    assert out.params['amplitude'].value < 10.0
    assert out.params['sigma1'].value > 1.0
    assert out.params['sigma1'].value < 5.0
    assert out.params['sigma2'].value > 0.3
    assert out.params['sigma2'].value < 2.5