Coverage for test/test_fitting.py: 100%

1import unittest

2from unittest.mock import MagicMock

3from pathlib import Path

4from collections import namedtuple

6import numpy as np

7import pandas as pd

8import pytest

9from pytest import fixture

10import nmrglue as ng

11from numpy.testing import assert_array_equal

12from lmfit import Model, Parameters

13from lmfit.model import ModelResult

15from peakipy.io import Pseudo3D

16from peakipy.fitting import (

17 FitDataModel,

18 validate_fit_data,

19 validate_fit_dataframe,

20 select_reference_planes_using_indices,

21 slice_peaks_from_data_using_mask,

22 select_planes_above_threshold_from_masked_data,

23 get_limits_for_axis_in_points,

24 deal_with_peaks_on_edge_of_spectrum,

25 estimate_amplitude,

26 make_mask,

27 make_mask_from_peak_cluster,

28 make_meshgrid,

29 get_params,

30 fix_params,

31 make_param_dict,

32 to_prefix,

33 make_models,

34 PeakLimits,

35 update_params,

36 make_masks_from_plane_data,

37 get_fit_peaks_result_validation_model,

38 FitPeaksResultRowPVPV,

39 FitPeaksResultRowVoigt,

40 FitPeaksResultRowGLPV,

41 filter_peak_clusters_by_max_cluster_size,

42 set_parameters_to_fix_during_fit,

43 unpack_fitted_parameters_for_lineshape,

44 get_default_lineshape_param_names,

45 split_parameter_sets_by_peak,

46 get_prefix_from_parameter_names,

47 create_parameter_dict,

48 perform_initial_lineshape_fit_on_cluster_of_peaks,

49 merge_unpacked_parameters_with_metadata,

50 add_vclist_to_df,

51 update_cluster_df_with_fit_statistics,

52 rename_columns_for_compatibility,

53 FitPeaksArgs,

54 FitPeaksInput,

55 FitResult,

56)

57from peakipy.lineshapes import Lineshape, pvoigt2d, pv_pv

60@pytest.fixture

61def fitdatamodel_dict():

62 return FitDataModel(

63 plane=1,

64 clustid=1,

65 assignment="assignment",

66 memcnt=1,

67 amp=10.0,

68 height=10.0,

69 center_x_ppm=0.0,

70 center_y_ppm=0.0,

71 fwhm_x_hz=10.0,

72 fwhm_y_hz=10.0,

73 lineshape="PV",

74 x_radius=0.04,

75 y_radius=0.4,

76 center_x=0.0,

77 center_y=0.0,

78 sigma_x=1.0,

79 sigma_y=1.0,

80 ).model_dump()

83def test_validate_fit_data_PVGL(fitdatamodel_dict):

84 fitdatamodel_dict.update(dict(fraction=0.5))

85 validate_fit_data(fitdatamodel_dict)

87 fitdatamodel_dict.update(dict(lineshape="G"))

88 validate_fit_data(fitdatamodel_dict)

90 fitdatamodel_dict.update(dict(lineshape="L"))

91 validate_fit_data(fitdatamodel_dict)

93 fitdatamodel_dict.update(

94 dict(lineshape="V", fraction=0.5, gamma_x=1.0, gamma_y=1.0)

95 )

96 validate_fit_data(fitdatamodel_dict)

98 fitdatamodel_dict.update(dict(lineshape="PVPV", fraction_x=0.5, fraction_y=1.0))

99 validate_fit_data(fitdatamodel_dict)

100

101

102def test_validate_fit_dataframe(fitdatamodel_dict):

103 fitdatamodel_dict.update(dict(fraction=0.5))

104 df = pd.DataFrame([fitdatamodel_dict] * 5)

105 validate_fit_dataframe(df)

106

107

108def test_select_reference_planes_using_indices():

109 data = np.zeros((6, 100, 200))

110 indices = []

111 np.testing.assert_array_equal(

112 select_reference_planes_using_indices(data, indices), data

113 )

114 indices = [1]

115 assert select_reference_planes_using_indices(data, indices).shape == (1, 100, 200)

116 indices = [1, -1]

117 assert select_reference_planes_using_indices(data, indices).shape == (2, 100, 200)

118

119

120def test_select_reference_planes_using_indices_min_index_error():

121 data = np.zeros((6, 100, 200))

122 indices = [-7]

123 with pytest.raises(IndexError):

124 select_reference_planes_using_indices(data, indices)

125

126

127def test_select_reference_planes_using_indices_max_index_error():

128 data = np.zeros((6, 100, 200))

129 indices = [6]

130 with pytest.raises(IndexError):

131 select_reference_planes_using_indices(data, indices)

132

133

134def test_slice_peaks_from_data_using_mask():

135 data = np.array(

136 [

137 np.array(

138 [

139 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],

140 [0, 0, 0, 0, 1, 1, 0, 0, 0, 0],

141 [0, 0, 0, 1, 2, 2, 1, 0, 0, 0],

142 [0, 0, 1, 2, 3, 3, 2, 1, 0, 0],

143 [0, 1, 2, 3, 4, 4, 3, 2, 1, 0],

144 [1, 2, 3, 4, 5, 5, 4, 3, 2, 1],

145 [0, 1, 2, 3, 4, 4, 3, 2, 1, 0],

146 [0, 0, 1, 2, 3, 3, 2, 1, 0, 0],

147 [0, 0, 0, 1, 2, 2, 1, 0, 0, 0],

148 [0, 0, 0, 0, 1, 1, 0, 0, 0, 0],

149 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],

150 ]

151 )

152 for i in range(5)

153 ]

154 )

155 mask = data[0] > 0

156 assert data.shape == (5, 11, 10)

157 assert mask.shape == (11, 10)

158 peak_slices = slice_peaks_from_data_using_mask(data, mask)

159 # array is flattened by application of mask

160 assert peak_slices.shape == (5, 50)

161

162

163def test_select_planes_above_threshold_from_masked_data():

164 peak_slices = np.array(

165 [

166 [1, 1, 1, 1, 1, 1],

167 [2, 2, 2, 2, 2, 2],

168 [-1, -1, -1, -1, -1, -1],

169 [-2, -2, -2, -2, -2, -2],

170 ]

171 )

172 assert peak_slices.shape == (4, 6)

173 threshold = -1

174 assert select_planes_above_threshold_from_masked_data(

175 peak_slices, threshold

176 ).shape == (

177 4,

178 6,

179 )

180 threshold = 2

181 assert_array_equal(

182 select_planes_above_threshold_from_masked_data(peak_slices, threshold),

183 peak_slices,

184 )

185 threshold = 1

186 assert select_planes_above_threshold_from_masked_data(

187 peak_slices, threshold

188 ).shape == (2, 6)

189

190 threshold = None

191 assert_array_equal(

192 select_planes_above_threshold_from_masked_data(peak_slices, threshold),

193 peak_slices,

194 )

195 threshold = 10

196 assert_array_equal(

197 select_planes_above_threshold_from_masked_data(peak_slices, threshold),

198 peak_slices,

199 )

200

201

202def test_make_param_dict():

203 selected_planes = [1, 2]

204 data = np.ones((4, 10, 5))

205 expected_shape = (2, 10, 5)

206 actual_shape = data[np.array(selected_planes)].shape

207 assert expected_shape == actual_shape

208

209

210def test_make_param_dict_sum():

211 data = np.ones((4, 10, 5))

212 expected_sum = 200

213 actual_sum = data.sum()

214 assert expected_sum == actual_sum

215

216

217def test_make_param_dict_selected():

218 selected_planes = [1, 2]

219 data = np.ones((4, 10, 5))

220 data = data[np.array(selected_planes)]

221 expected_sum = 100

222 actual_sum = data.sum()

223 assert expected_sum == actual_sum

224

225

226def test_update_params_normal_case():

227 params = Parameters()

228 params.add("center_x", value=0)

229 params.add("center_y", value=0)

230 params.add("sigma", value=1)

231 params.add("gamma", value=1)

232 params.add("fraction", value=0.5)

233

234 param_dict = {

235 "center_x": 10,

236 "center_y": 20,

237 "sigma": 2,

238 "gamma": 3,

239 "fraction": 0.8,

240 }

241

242 xy_bounds = (5, 5)

243

244 update_params(params, param_dict, Lineshape.PV, xy_bounds)

245

246 assert params["center_x"].value == 10

247 assert params["center_y"].value == 20

248 assert params["sigma"].value == 2

249 assert params["gamma"].value == 3

250 assert params["fraction"].value == 0.8

251 assert params["center_x"].min == 5

252 assert params["center_x"].max == 15

253 assert params["center_y"].min == 15

254 assert params["center_y"].max == 25

255 assert params["sigma"].min == 0.0

256 assert params["sigma"].max == 1e4

257 assert params["gamma"].min == 0.0

258 assert params["gamma"].max == 1e4

259 assert params["fraction"].min == 0.0

260 assert params["fraction"].max == 1.0

261 assert params["fraction"].vary is True

262

263

264def test_update_params_lineshape_G():

265 params = Parameters()

266 params.add("fraction", value=0.5)

267

268 param_dict = {"fraction": 0.7}

269

270 update_params(params, param_dict, Lineshape.G)

271

272 assert params["fraction"].value == 0.7

273 assert params["fraction"].min == 0.0

274 assert params["fraction"].max == 1.0

275 assert params["fraction"].vary is False

276

277

278def test_update_params_lineshape_L():

279 params = Parameters()

280 params.add("fraction", value=0.5)

281

282 param_dict = {"fraction": 0.7}

283

284 update_params(params, param_dict, Lineshape.L)

285

286 assert params["fraction"].value == 0.7

287 assert params["fraction"].min == 0.0

288 assert params["fraction"].max == 1.0

289 assert params["fraction"].vary is False

290

291

292def test_update_params_lineshape_PV_PV():

293 params = Parameters()

294 params.add("fraction", value=0.5)

295

296 param_dict = {"fraction": 0.7}

297

298 update_params(params, param_dict, Lineshape.PV_PV)

299

300 assert params["fraction"].value == 0.7

301 assert params["fraction"].min == 0.0

302 assert params["fraction"].max == 1.0

303 assert params["fraction"].vary is True

304

305

306def test_update_params_no_bounds():

307 params = Parameters()

308 params.add("center_x", value=0)

309 params.add("center_y", value=0)

310

311 param_dict = {

312 "center_x": 10,

313 "center_y": 20,

314 }

315

316 update_params(params, param_dict, Lineshape.PV, None)

317

318 assert params["center_x"].value == 10

319 assert params["center_y"].value == 20

320 assert params["center_x"].min == -np.inf

321 assert params["center_x"].max == np.inf

322 assert params["center_y"].min == -np.inf

323 assert params["center_y"].max == np.inf

324

325

326def test_peak_limits_normal_case():

327 peak = pd.DataFrame({"X_AXIS": [5], "Y_AXIS": [5], "XW": [2], "YW": [2]}).iloc[0]

328 data = np.zeros((10, 10))

329 pl = PeakLimits(peak, data)

330 assert pl.min_x == 3

331 assert pl.max_x == 8

332 assert pl.min_y == 3

333 assert pl.max_y == 8

334

335

336def test_peak_limits_at_edge():

337 peak = pd.DataFrame({"X_AXIS": [1], "Y_AXIS": [1], "XW": [2], "YW": [2]}).iloc[0]

338 data = np.zeros((10, 10))

339 pl = PeakLimits(peak, data)

340 assert pl.min_x == 0

341 assert pl.max_x == 4

342 assert pl.min_y == 0

343 assert pl.max_y == 4

344

345

346def test_peak_limits_exceeding_bounds():

347 peak = pd.DataFrame({"X_AXIS": [9], "Y_AXIS": [9], "XW": [2], "YW": [2]}).iloc[0]

348 data = np.zeros((10, 10))

349 pl = PeakLimits(peak, data)

350 assert pl.min_x == 7

351 assert pl.max_x == 10

352 assert pl.min_y == 7

353 assert pl.max_y == 10

354

355

356def test_peak_limits_small_data():

357 peak = pd.DataFrame({"X_AXIS": [2], "Y_AXIS": [2], "XW": [5], "YW": [5]}).iloc[0]

358 data = np.zeros((5, 5))

359 pl = PeakLimits(peak, data)

360 assert pl.min_x == 0

361 assert pl.max_x == 5

362 assert pl.min_y == 0

363 assert pl.max_y == 5

364

365

366def test_peak_limits_assertion_error():

367 peak = pd.DataFrame({"X_AXIS": [11], "Y_AXIS": [11], "XW": [2], "YW": [2]}).iloc[0]

368 data = np.zeros((10, 10))

369 with pytest.raises(AssertionError):

370 pl = PeakLimits(peak, data)

371

372

373def test_estimate_amplitude():

374 peak = namedtuple("peak", ["X_AXIS", "XW", "Y_AXIS", "YW"])

375 p = peak(5, 2, 3, 2)

376 data = np.ones((20, 10))

377 expected_result = 25

378 actual_result = estimate_amplitude(p, data)

379 assert expected_result == actual_result

380

381

382def test_estimate_amplitude_invalid_indices():

383 peak = namedtuple("peak", ["X_AXIS", "XW", "Y_AXIS", "YW"])

384 p = peak(1, 2, 3, 2)

385 data = np.ones((20, 10))

386 expected_result = 20

387 actual_result = estimate_amplitude(p, data)

388 assert expected_result == actual_result

389

390

391def test_make_mask_from_peak_cluster():

392 data = np.ones((10, 10))

393 group = pd.DataFrame(

394 {"X_AXISf": [3, 6], "Y_AXISf": [3, 6], "X_RADIUS": [2, 3], "Y_RADIUS": [2, 3]}

395 )

396 mask, peak = make_mask_from_peak_cluster(group, data)

397 expected_mask = np.array(

398 [

399 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],

400 [0, 0, 0, 1, 0, 0, 0, 0, 0, 0],

401 [0, 0, 1, 1, 1, 0, 0, 0, 0, 0],

402 [0, 1, 1, 1, 1, 1, 1, 0, 0, 0],

403 [0, 0, 1, 1, 1, 1, 1, 1, 1, 0],

404 [0, 0, 0, 1, 1, 1, 1, 1, 1, 0],

405 [0, 0, 0, 1, 1, 1, 1, 1, 1, 1],

406 [0, 0, 0, 0, 1, 1, 1, 1, 1, 0],

407 [0, 0, 0, 0, 1, 1, 1, 1, 1, 0],

408 [0, 0, 0, 0, 0, 0, 1, 0, 0, 0],

409 ],

410 dtype=bool,

411 )

412 assert_array_equal(expected_mask, mask)

413

414

415# get_limits_for_axis_in_points

416def test_positive_points():

417 group_axis_points = np.array([1, 2, 3, 4, 5])

418 mask_radius_in_points = 2

419 expected = (8, -1) # ceil(5+1+1), floor(1-1)

420 assert (

421 get_limits_for_axis_in_points(group_axis_points, mask_radius_in_points)

422 == expected

423 )

424

425

426def test_single_point():

427 group_axis_points = np.array([5])

428 mask_radius_in_points = 3

429 expected = (9, 2)

430 assert (

431 get_limits_for_axis_in_points(group_axis_points, mask_radius_in_points)

432 == expected

433 )

434

435

436def test_no_radius():

437 group_axis_points = np.array([1, 2, 3])

438 mask_radius_in_points = 0

439 expected = (4, 1)

440 assert (

441 get_limits_for_axis_in_points(group_axis_points, mask_radius_in_points)

442 == expected

443 )

444

445

446# deal_with_peaks_on_edge_of_spectrum

447def test_min_y_less_than_zero():

448 assert deal_with_peaks_on_edge_of_spectrum((100, 200), 50, 30, 10, -10) == (

449 50,

450 30,

451 10,

452 0,

453 )

454

455

456def test_min_x_less_than_zero():

457 assert deal_with_peaks_on_edge_of_spectrum((100, 200), 50, -5, 70, 20) == (

458 50,

459 0,

460 70,

461 20,

462 )

463

464

465def test_max_y_exceeds_data_shape():

466 assert deal_with_peaks_on_edge_of_spectrum((100, 200), 50, 30, 110, 20) == (

467 50,

468 30,

469 100,

470 20,

471 )

472

473

474def test_max_x_exceeds_data_shape():

475 assert deal_with_peaks_on_edge_of_spectrum((100, 200), 250, 30, 70, 20) == (

476 200,

477 30,

478 70,

479 20,

480 )

481

482

483def test_values_within_range():

484 assert deal_with_peaks_on_edge_of_spectrum((100, 200), 50, 30, 70, 20) == (

485 50,

486 30,

487 70,

488 20,

489 )

490

491

492def test_all_edge_cases():

493 assert deal_with_peaks_on_edge_of_spectrum((100, 200), 250, -5, 110, -10) == (

494 200,

495 0,

496 100,

497 0,

498 )

499

500

501def test_make_meshgrid():

502 data_shape = (4, 5)

503 expected_x = np.array(

504 [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]

505 )

506 expected_y = np.array(

507 [[0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [2, 2, 2, 2, 2], [3, 3, 3, 3, 3]]

508 )

509 XY = make_meshgrid(data_shape)

510 np.testing.assert_array_equal(XY[0], expected_x)

511 np.testing.assert_array_equal(XY[1], expected_y)

512

513

514class TestCoreFunctions(unittest.TestCase):

515 test_directory = Path(__file__).parent

516 test_directory = "./test"

517

518 def test_make_mask(self):

519 data = np.ones((10, 10))

520 c_x = 5

521 c_y = 5

522 r_x = 3

523 r_y = 2

524

525 expected_result = np.array(

526 [

527 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],

528 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],

529 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],

530 [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],

531 [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0],

532 [0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0],

533 [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0],

534 [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],

535 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],

536 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],

537 ]

538 )

539

540 result = np.array(make_mask(data, c_x, c_y, r_x, r_y), dtype=int)

541 test = result - expected_result

542 # print(test)

543 # print(test.sum())

544 # print(result)

545 self.assertEqual(test.sum(), 0)

546

547 def test_make_mask_2(self):

548 data = np.ones((10, 10))

549 c_x = 5

550 c_y = 8

551 r_x = 3

552 r_y = 2

553

554 expected_result = np.array(

555 [

556 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],

557 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],

558 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],

559 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],

560 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],

561 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],

562 [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],

563 [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0],

564 [0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0],

565 [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0],

566 ]

567 )

568

569 result = np.array(make_mask(data, c_x, c_y, r_x, r_y), dtype=int)

570 test = result - expected_result

571 # print(test)

572 # print(test.sum())

573 # print(result)

574 self.assertEqual(test.sum(), 0)

575

576 def test_fix_params(self):

577 mod = Model(pvoigt2d)

578 pars = mod.make_params()

579 to_fix = ["center", "sigma", "fraction"]

580 fix_params(pars, to_fix)

581

582 self.assertEqual(pars["center_x"].vary, False)

583 self.assertEqual(pars["center_y"].vary, False)

584 self.assertEqual(pars["sigma_x"].vary, False)

585 self.assertEqual(pars["sigma_y"].vary, False)

586 self.assertEqual(pars["fraction"].vary, False)

587

588 def test_get_params(self):

589 mod = Model(pvoigt2d, prefix="p1_")

590 pars = mod.make_params(p1_center_x=20.0, p1_center_y=30.0)

591 pars["p1_center_x"].stderr = 1.0

592 pars["p1_center_y"].stderr = 2.0

593 ps, ps_err, names, prefixes = get_params(pars, "center")

594 # get index of values

595 cen_x = names.index("p1_center_x")

596 cen_y = names.index("p1_center_y")

597

598 self.assertEqual(ps[cen_x], 20.0)

599 self.assertEqual(ps[cen_y], 30.0)

600 self.assertEqual(ps_err[cen_x], 1.0)

601 self.assertEqual(ps_err[cen_y], 2.0)

602 self.assertEqual(prefixes[cen_y], "p1_")

603

604 def test_make_param_dict(self):

605 peaks = pd.DataFrame(

606 {

607 "ASS": ["one", "two", "three"],

608 "X_AXISf": [5.0, 10.0, 15.0],

609 "X_AXIS": [5, 10, 15],

610 "Y_AXISf": [15.0, 10.0, 5.0],

611 "Y_AXIS": [15, 10, 5],

612 "XW": [2.5, 2.5, 2.5],

613 "YW": [2.5, 2.5, 2.5],

614 }

615 )

616 data = np.ones((20, 20))

617

618 for ls, frac in zip([Lineshape.PV, Lineshape.G, Lineshape.L], [0.5, 0.0, 1.0]):

619 params = make_param_dict(peaks, data, ls)

620 self.assertEqual(params["_one_fraction"], frac)

621 self.assertEqual(params["_two_fraction"], frac)

622 self.assertEqual(params["_three_fraction"], frac)

623

624 self.assertEqual(params["_one_center_x"], 5.0)

625 self.assertEqual(params["_two_center_x"], 10.0)

626 self.assertEqual(params["_two_sigma_x"], 1.25)

627 self.assertEqual(params["_two_sigma_y"], 1.25)

628

629 voigt_params = make_param_dict(peaks, data, Lineshape.V)

630 self.assertEqual(

631 voigt_params["_one_sigma_x"], 2.5 / (2.0 * np.sqrt(2.0 * np.log(2)))

632 )

633 self.assertEqual(voigt_params["_one_gamma_x"], 2.5 / 2.0)

634

635 def test_to_prefix(self):

636 names = [

637 (1, "_1_"),

638 (1.0, "_1_0_"),

639 (" one", "_one_"),

640 (" one/two", "_oneortwo_"),

641 (" one?two", "_onemaybetwo_"),

642 (r" [{one?two\}][", "___onemaybetwo____"),

643 ("hel'lo", "_hel_lo_"),

644 ]

645 for test, expect in names:

646 prefix = to_prefix(test)

647 # print(prefix)

648 self.assertEqual(prefix, expect)

649

650 def test_make_models(self):

651 peaks = pd.DataFrame(

652 {

653 "ASS": ["one", "two", "three"],

654 "X_AXISf": [5.0, 10.0, 15.0],

655 "X_AXIS": [5, 10, 15],

656 "Y_AXISf": [15.0, 10.0, 5.0],

657 "Y_AXIS": [15, 10, 5],

658 "XW": [2.5, 2.5, 2.5],

659 "YW": [2.5, 2.5, 2.5],

660 "CLUSTID": [1, 1, 1],

661 }

662 )

663

664 group = peaks.groupby("CLUSTID")

665

666 data = np.ones((20, 20))

667

668 lineshapes = [Lineshape.PV, Lineshape.L, Lineshape.G, Lineshape.PV_PV]

669

670 for lineshape in lineshapes:

671 match lineshape:

672 case lineshape.PV:

673 mod, p_guess = make_models(pvoigt2d, peaks, data, lineshape)

674 self.assertEqual(p_guess["_one_fraction"].vary, True)

675 self.assertEqual(p_guess["_one_fraction"].value, 0.5)

676

677 case lineshape.G:

678 mod, p_guess = make_models(pvoigt2d, peaks, data, lineshape)

679 self.assertEqual(p_guess["_one_fraction"].vary, False)

680 self.assertEqual(p_guess["_one_fraction"].value, 0.0)

681

682 case lineshape.L:

683 mod, p_guess = make_models(pvoigt2d, peaks, data, lineshape)

684 self.assertEqual(p_guess["_one_fraction"].vary, False)

685 self.assertEqual(p_guess["_one_fraction"].value, 1.0)

686

687 case lineshape.PV_PV:

688 mod, p_guess = make_models(pv_pv, peaks, data, lineshape)

689 self.assertEqual(p_guess["_one_fraction_x"].vary, True)

690 self.assertEqual(p_guess["_one_fraction_x"].value, 0.5)

691 self.assertEqual(p_guess["_one_fraction_y"].vary, True)

692 self.assertEqual(p_guess["_one_fraction_y"].value, 0.5)

693

694 def test_Pseudo3D(self):

695 datasets = [

696 (f"{self.test_directory}/test_protein_L/test1.ft2", [0, 1, 2]),

697 (f"{self.test_directory}/test_protein_L/test_tp.ft2", [2, 1, 0]),

698 (f"{self.test_directory}/test_protein_L/test_tp2.ft2", [1, 2, 0]),

699 ]

700

701 # expected shape

702 data_shape = (4, 256, 546)

703 test_nu = 1

704 for dataset, dims in datasets:

705 with self.subTest(i=test_nu):

706 dic, data = ng.pipe.read(dataset)

707 pseudo3D = Pseudo3D(dic, data, dims)

708 self.assertEqual(dims, pseudo3D.dims)

709 self.assertEqual(pseudo3D.data.shape, data_shape)

710 self.assertEqual(pseudo3D.f1_label, "15N")

711 self.assertEqual(pseudo3D.f2_label, "HN")

712 self.assertEqual(pseudo3D.dims, dims)

713 self.assertEqual(pseudo3D.f1_size, 256)

714 self.assertEqual(pseudo3D.f2_size, 546)

715 test_nu += 1

716

717

718def test_get_fit_peaks_result_validation_model_PVPV():

719 validation_model = get_fit_peaks_result_validation_model(Lineshape.PV_PV)

720 assert validation_model == FitPeaksResultRowPVPV

721

722

723def test_get_fit_peaks_result_validation_model_G():

724 validation_model = get_fit_peaks_result_validation_model(Lineshape.G)

725 assert validation_model == FitPeaksResultRowGLPV

726

727

728def test_get_fit_peaks_result_validation_model_L():

729 validation_model = get_fit_peaks_result_validation_model(Lineshape.L)

730 assert validation_model == FitPeaksResultRowGLPV

731

732

733def test_get_fit_peaks_result_validation_model_PV():

734 validation_model = get_fit_peaks_result_validation_model(Lineshape.PV)

735 assert validation_model == FitPeaksResultRowGLPV

736

737

738def test_get_fit_peaks_result_validation_model_V():

739 validation_model = get_fit_peaks_result_validation_model(Lineshape.V)

740 assert validation_model == FitPeaksResultRowVoigt

741

742

743def test_filter_groups_by_max_cluster_size():

744 groups = pd.DataFrame(

745 dict(

746 col1=[1, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 4, 5, 6, 7],

747 col2=[1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7],

748 )

749 ).groupby("col1")

750 max_cluster_size = 3

751 filtered_groups = filter_peak_clusters_by_max_cluster_size(groups, max_cluster_size)

752 filtered_group_names = filtered_groups.col1.unique()

753 expected_group_names = np.array([3, 4, 5, 6, 7])

754 np.testing.assert_array_equal(filtered_group_names, expected_group_names)

755

756

757def test_set_parameters_to_fix_during_fit():

758 parameter_set = Parameters()

759 parameter_set.add("test1", vary=True)

760 modified_parameter_set, float_str = set_parameters_to_fix_during_fit(

761 parameter_set, ["test1"]

762 )

763 assert modified_parameter_set["test1"].vary == False

764

765

766@fixture

767def parameters_set_with_two_variables():

768 parameter_set = Parameters()

769 parameter_set.add("prefix1_test1", vary=True)

770 parameter_set.add("prefix1_test2", vary=True)

771 return parameter_set

772

773

774def test_set_parameters_to_fix_during_fit_2(parameters_set_with_two_variables):

775 modified_parameter_set, float_str = set_parameters_to_fix_during_fit(

776 parameters_set_with_two_variables, ["prefix1_test1", "prefix1_test2"]

777 )

778 assert (

779 modified_parameter_set["prefix1_test2"].vary

780 == modified_parameter_set["prefix1_test1"].vary

781 == False

782 )

783

784

785def test_set_parameters_to_fix_during_fit_3():

786 parameter_set = Parameters()

787 parameter_set.add("test1", vary=True)

788 parameter_set.add("test2", vary=True)

789 modified_parameter_set, float_str = set_parameters_to_fix_during_fit(

790 parameter_set, ["test2"]

791 )

792 assert (

793 modified_parameter_set["test1"].vary

794 != modified_parameter_set["test2"].vary

795 == False

796 )

797

798

799def test_set_parameters_to_fix_during_fit_None():

800 parameter_set = Parameters()

801 parameter_set.add("test1", vary=True)

802 parameter_set.add("test2", vary=True)

803 modified_parameter_set, float_str = set_parameters_to_fix_during_fit(

804 parameter_set, None

805 )

806 assert (

807 modified_parameter_set["test1"].vary

808 == modified_parameter_set["test2"].vary

809 == True

810 )

811

812

813def test_set_parameters_to_fix_during_fit_None_str():

814 parameter_set = Parameters()

815 parameter_set.add("test1", vary=True)

816 parameter_set.add("test2", vary=True)

817 modified_parameter_set, float_str = set_parameters_to_fix_during_fit(

818 parameter_set, ["None"]

819 )

820 assert (

821 modified_parameter_set["test1"].vary

822 == modified_parameter_set["test2"].vary

823 == True

824 )

825

826

827def test_update_cluster_df_with_fit_statistics():

828 result = ModelResult(Model(pvoigt2d), None, None)

829 result.aic = None

830 result.bic = None

831 data = [

832 dict(

833 chisqr=None,

834 redchi=None,

835 residual_sum=None,

836 aic=None,

837 bic=None,

838 nfev=0,

839 ndata=0,

840 )

841 ]

842 expected_cluster_df = pd.DataFrame(data)

843 actual_cluster_df = update_cluster_df_with_fit_statistics(

844 expected_cluster_df, result

845 )

846 pd.testing.assert_frame_equal(actual_cluster_df, expected_cluster_df)

847

848

849def test_rename_columns_for_compatibility():

850 df = pd.DataFrame(

851 [

852 dict(

853 amplitude=1,

854 amplitude_stderr=1,

855 X_AXIS=1,

856 Y_AXIS=1,

857 ASS="None",

858 MEMCNT=1,

859 X_RADIUS=1,

860 Y_RADIUS=1,

861 )

862 ]

863 )

864 expected_columns = [

865 "amp",

866 "amp_err",

867 "init_center_x",

868 "init_center_y",

869 "assignment",

870 "memcnt",

871 "x_radius",

872 "y_radius",

873 ]

874 actual_columns = rename_columns_for_compatibility(df).columns

875 assert all([i == j for i, j in zip(actual_columns, expected_columns)])

876

877

878def test_get_default_param_names_pseudo_voigt():

879 assert get_default_lineshape_param_names(Lineshape.PV) == [

880 "amplitude",

881 "center_x",

882 "center_y",

883 "sigma_x",

884 "sigma_y",

885 "fraction",

886 ]

887

888

889def test_get_default_param_names_gaussian():

890 assert get_default_lineshape_param_names(Lineshape.G) == [

891 "amplitude",

892 "center_x",

893 "center_y",

894 "sigma_x",

895 "sigma_y",

896 "fraction",

897 ]

898

899

900def test_get_default_param_names_lorentzian():

901 assert get_default_lineshape_param_names(Lineshape.L) == [

902 "amplitude",

903 "center_x",

904 "center_y",

905 "sigma_x",

906 "sigma_y",

907 "fraction",

908 ]

909

910

911def test_get_default_param_names_pv_pv():

912 assert get_default_lineshape_param_names(Lineshape.PV_PV) == [

913 "amplitude",

914 "center_x",

915 "center_y",

916 "sigma_x",

917 "sigma_y",

918 "fraction_x",

919 "fraction_y",

920 ]

921

922

923def test_get_default_param_names_voigt():

924 assert get_default_lineshape_param_names(Lineshape.V) == [

925 "amplitude",

926 "center_x",

927 "center_y",

928 "sigma_x",

929 "sigma_y",

930 "gamma_x",

931 "gamma_y",

932 "fraction",

933 ]

934

935

936def test_split_parameter_sets_by_peak(default_pseudo_voigt_parameter_names):

937 # the second element of each tuple actually contains an

938 # lmfit.Parameter object

939 params = [

940 ("p1_amplitude", "amplitude"),

941 ("p1_center_x", "center_x"),

942 ("p1_center_y", "center_y"),

943 ("p1_sigma_x", "sigma_x"),

944 ("p1_sigma_y", "sigma_y"),

945 ("p1_fraction", "fraction"),

946 ("p2_amplitude", "amplitude"),

947 ("p2_center_x", "center_x"),

948 ("p2_center_y", "center_y"),

949 ("p2_sigma_x", "sigma_x"),

950 ("p2_sigma_y", "sigma_y"),

951 ("p2_fraction", "fraction"),

952 ("p3_amplitude", "amplitude"),

953 ("p3_center_x", "center_x"),

954 ("p3_center_y", "center_y"),

955 ("p3_sigma_x", "sigma_x"),

956 ("p3_sigma_y", "sigma_y"),

957 ("p3_fraction", "fraction"),

958 ]

959 expected_result = [

960 [

961 ("p1_amplitude", "amplitude"),

962 ("p1_center_x", "center_x"),

963 ("p1_center_y", "center_y"),

964 ("p1_sigma_x", "sigma_x"),

965 ("p1_sigma_y", "sigma_y"),

966 ("p1_fraction", "fraction"),

967 ],

968 [

969 ("p2_amplitude", "amplitude"),

970 ("p2_center_x", "center_x"),

971 ("p2_center_y", "center_y"),

972 ("p2_sigma_x", "sigma_x"),

973 ("p2_sigma_y", "sigma_y"),

974 ("p2_fraction", "fraction"),

975 ],

976 [

977 ("p3_amplitude", "amplitude"),

978 ("p3_center_x", "center_x"),

979 ("p3_center_y", "center_y"),

980 ("p3_sigma_x", "sigma_x"),

981 ("p3_sigma_y", "sigma_y"),

982 ("p3_fraction", "fraction"),

983 ],

984 ]

985 expected_result_parameter_names = [[j[0] for j in i] for i in expected_result]

986 split_parameter_names = [

987 [j[0] for j in i]

988 for i in split_parameter_sets_by_peak(

989 default_pseudo_voigt_parameter_names, params

990 )

991 ]

992 assert split_parameter_names == expected_result_parameter_names

993

994

995@fixture

996def default_pseudo_voigt_parameter_names():

997 return Model(pvoigt2d).param_names

998

999

1000def test_get_prefix_from_parameter_names(default_pseudo_voigt_parameter_names):

1001 parameter_items_with_prefixes = [

1002 ("p1_amplitude", "amplitude"),

1003 ("p1_center_x", "center_x"),

1004 ("p1_center_y", "center_y"),

1005 ("p1_sigma_x", "sigma_x"),

1006 ("p1_sigma_y", "sigma_y"),

1007 ("p1_fraction", "fraction"),

1008 ]

1009 expected_result = "p1_"

1010 actual_result = get_prefix_from_parameter_names(

1011 default_pseudo_voigt_parameter_names, parameter_items_with_prefixes

1012 )

1013 assert expected_result == actual_result

1014

1015

1016@fixture

1017def pseudo_voigt_model_result():

1018 m1 = Model(pvoigt2d, prefix="p1_")

1019 m2 = Model(pvoigt2d, prefix="p2_")

1020 model = m1 + m2

1021 params = model.make_params()

1022 model_result = ModelResult(model, params)

1023 return model_result

1024

1025

1026def test_create_parameter_dict(pseudo_voigt_model_result):

1027 prefix = "p1_"

1028 params = list(pseudo_voigt_model_result.params.items())[:6]

1029 expected_result = dict(

1030 prefix="p1_",

1031 amplitude=1.0,

1032 amplitude_stderr=None,

1033 center_x=0.5,

1034 center_x_stderr=None,

1035 center_y=0.5,

1036 center_y_stderr=None,

1037 sigma_x=1.0,

1038 sigma_x_stderr=None,

1039 sigma_y=1.0,

1040 sigma_y_stderr=None,

1041 fraction=0.5,

1042 fraction_stderr=None,

1043 )

1044 actual_result = create_parameter_dict(prefix, params)

1045 assert expected_result == actual_result

1046

1047

1048def test_unpack_fitted_parameters_for_lineshape_PV(pseudo_voigt_model_result):

1049 expected_params = [

1050 dict(

1051 prefix="p1_",

1052 plane=0,

1053 amplitude=1.0,

1054 amplitude_stderr=None,

1055 center_x=0.5,

1056 center_x_stderr=None,

1057 center_y=0.5,

1058 center_y_stderr=None,

1059 sigma_x=1.0,

1060 sigma_x_stderr=None,

1061 sigma_y=1.0,

1062 sigma_y_stderr=None,

1063 fraction=0.5,

1064 fraction_stderr=None,

1065 ),

1066 dict(

1067 prefix="p2_",

1068 plane=0,

1069 amplitude=1.0,

1070 amplitude_stderr=None,

1071 center_x=0.5,

1072 center_x_stderr=None,

1073 center_y=0.5,

1074 center_y_stderr=None,

1075 sigma_x=1.0,

1076 sigma_x_stderr=None,

1077 sigma_y=1.0,

1078 sigma_y_stderr=None,

1079 fraction=0.5,

1080 fraction_stderr=None,

1081 ),

1082 ]

1083 unpacked_params = unpack_fitted_parameters_for_lineshape(

1084 Lineshape.PV, list(pseudo_voigt_model_result.params.items()), plane_number=0

1085 )

1086 assert expected_params == unpacked_params

1087

1088

1089def test_merge_unpacked_parameters_with_metadata():

1090 cluster_fit_df = pd.DataFrame(

1091 dict(

1092 plane=[0, 1, 2, 3, 0, 1, 2, 3],

1093 prefix=["_p1_", "_p1_", "_p1_", "_p1_", "_p2_", "_p2_", "_p2_", "_p2_"],

1094 )

1095 )

1096 peak_df = pd.DataFrame(dict(ASS=["p1", "p2"], data=["p1_data", "p2_data"]))

1097 expected_result = pd.DataFrame(

1098 dict(

1099 plane=[0, 1, 2, 3, 0, 1, 2, 3],

1100 prefix=["_p1_", "_p1_", "_p1_", "_p1_", "_p2_", "_p2_", "_p2_", "_p2_"],

1101 ASS=["p1", "p1", "p1", "p1", "p2", "p2", "p2", "p2"],

1102 data=[

1103 "p1_data",

1104 "p1_data",

1105 "p1_data",

1106 "p1_data",

1107 "p2_data",

1108 "p2_data",

1109 "p2_data",

1110 "p2_data",

1111 ],

1112 )

1113 )

1114 actual_result = merge_unpacked_parameters_with_metadata(cluster_fit_df, peak_df)

1115 assert expected_result.equals(actual_result)

1116

1117

1118def test_add_vclist_to_df():

1119 args = FitPeaksArgs(

1120 noise=0, uc_dics={}, lineshape=Lineshape.PV, vclist_data=np.array([1, 2, 3])

1121 )

1122 fit_peaks_input = FitPeaksInput(

1123 args=args, data=None, config=None, plane_numbers=None

1124 )

1125 df = pd.DataFrame(dict(plane=[0, 1, 2]))

1126 expected_df = pd.DataFrame(dict(plane=[0, 1, 2], vclist=[1, 2, 3]))

1127 actual_df = add_vclist_to_df(fit_peaks_input, df)

1128 assert actual_df.equals(expected_df)

1129

1130

1131def test_add_vclist_to_df_plane_order():

1132 args = FitPeaksArgs(

1133 noise=0, uc_dics={}, lineshape=Lineshape.PV, vclist_data=np.array([1, 2, 3])

1134 )

1135 fit_peaks_input = FitPeaksInput(

1136 args=args, data=None, config=None, plane_numbers=None

1137 )

1138 df = pd.DataFrame(dict(plane=[2, 1, 0]))

1139 expected_df = pd.DataFrame(dict(plane=[2, 1, 0], vclist=[3, 2, 1]))

1140 actual_df = add_vclist_to_df(fit_peaks_input, df)

1141 assert actual_df.equals(expected_df)

1142

1143

1144# def test_perform_initial_lineshape_fit_on_cluster_of_peaks(pseudo_voigt_model_result):

1145# expected_result = pseudo_voigt_model_result

1146# actual_result = perform_initial_lineshape_fit_on_cluster_of_peaks()

1147# assert expected_result == actual_result

1148# Mock FitPeakClusterInput class for testing purposes

1149class MockFitPeakClusterInput:

1150 def __init__(

1151 self,

1152 mod,

1153 peak_slices,

1154 XY_slices,

1155 p_guess,

1156 weights,

1157 fit_method,

1158 mask,

1159 XY,

1160 first_plane_data,

1161 last_peak,

1162 group,

1163 min_x,

1164 min_y,

1165 max_x,

1166 max_y,

1167 verbose,

1168 uc_dics,

1169 ):

1170 self.mod = mod

1171 self.peak_slices = peak_slices

1172 self.XY_slices = XY_slices

1173 self.p_guess = p_guess

1174 self.weights = weights

1175 self.fit_method = fit_method

1176 self.mask = mask

1177 self.XY = XY

1178 self.first_plane_data = first_plane_data

1179 self.last_peak = last_peak

1180 self.group = group

1181 self.min_x = min_x

1182 self.min_y = min_y

1183 self.max_x = max_x

1184 self.max_y = max_y

1185 self.verbose = verbose

1186 self.uc_dics = uc_dics

1187

1188

1189@pytest.fixture

1190def fit_peak_cluster_input():

1191 mod = MagicMock()

1192 mod.fit = MagicMock(

1193 return_value=MagicMock(

1194 params="params", fit_report=MagicMock(return_value="fit_report")

1195 )

1196 )

1197 mod.eval = MagicMock(return_value=np.array([2.0, 1.0, 2.0]))

1198

1199 return MockFitPeakClusterInput(

1200 mod=mod,

1201 peak_slices="peak_slices",

1202 XY_slices="XY_slices",

1203 p_guess="p_guess",

1204 weights="weights",

1205 fit_method="fit_method",

1206 mask=np.array([True, False, True]),

1207 XY=(np.array([0, 1, 2]), np.array([0, 1, 2])),

1208 first_plane_data=np.array([2.0, 1.0, 2.0]),

1209 last_peak="last_peak",

1210 group="group",

1211 min_x="min_x",

1212 min_y="min_y",

1213 max_x="max_x",

1214 max_y="max_y",

1215 verbose=True,

1216 uc_dics="uc_dics",

1217 )

1218

1219

1220def test_perform_initial_lineshape_fit_on_cluster_of_peaks(fit_peak_cluster_input):

1221

1222 result = perform_initial_lineshape_fit_on_cluster_of_peaks(fit_peak_cluster_input)

1223

1224 # Check if result is an instance of FitResult

1225 assert isinstance(result, FitResult)

1226

1227 # Verify returned values

1228 assert result.out.params == "params"

1229 np.testing.assert_array_equal(result.mask, np.array([True, False, True]))

1230 assert result.fit_str == ""

1231 assert result.log == ""

1232 assert result.group == "group"

1233 assert result.uc_dics == "uc_dics"

1234 assert result.min_x == "min_x"

1235 assert result.min_y == "min_y"

1236 assert result.max_x == "max_x"

1237 assert result.max_y == "max_y"

1238 np.testing.assert_array_equal(result.X, np.array([0, 1, 2]))

1239 np.testing.assert_array_equal(result.Y, np.array([0, 1, 2]))

1240 np.testing.assert_array_equal(result.Z, np.array([2.0, np.nan, 2.0]))

1241 np.testing.assert_array_equal(result.Z_sim, np.array([2.0, np.nan, 2.0]))

1242 assert result.peak_slices == "peak_slices"

1243 assert result.XY_slices == "XY_slices"

1244 assert result.weights == "weights"

1245 assert result.mod == fit_peak_cluster_input.mod

1246

1247 # Check if mod.fit and mod.eval were called with correct arguments

1248 fit_peak_cluster_input.mod.fit.assert_called_once_with(

1249 "peak_slices",

1250 XY="XY_slices",

1251 params="p_guess",

1252 weights="weights",

1253 method="fit_method",

1254 )

1255 # fit_peak_cluster_input.mod.eval.assert_called_once_with(

1256 # XY=(np.array([0,1,2]), np.array([0,1,2])), params='params'

1257 # )