234 lines
7.6 KiB
Python
234 lines
7.6 KiB
Python
"""
|
|
@package tests.test_project
|
|
unit tests for pmsco.project.
|
|
|
|
the purpose of these tests is to help debugging the code.
|
|
|
|
to run the tests, change to the directory which contains the tests directory, and execute =nosetests=.
|
|
|
|
@pre nose and mock must be installed.
|
|
|
|
@author Matthias Muntwiler, matthias.muntwiler@psi.ch
|
|
|
|
@copyright (c) 2015-21 by Paul Scherrer Institut @n
|
|
Licensed under the Apache License, Version 2.0 (the "License"); @n
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
"""
|
|
|
|
import mock
|
|
import numpy as np
|
|
import os
|
|
from pathlib import Path
|
|
import unittest
|
|
|
|
import pmsco.data as data
|
|
import pmsco.dispatch as dispatch
|
|
import pmsco.project as project
|
|
|
|
|
|
class TestModelSpace(unittest.TestCase):
|
|
def setUp(self):
|
|
self.d1 = {
|
|
"A": {"start": 2.1, "min": 2.0, "max": 3.0, "step": 0.05},
|
|
"B": {"start": 15.0, "min": 0.0, "max": 30.0, "step": 1.0}}
|
|
self.d2 = {
|
|
"C": {"start": 22.0, "min": 15.0, "max": 25.0, "step": 1.0},
|
|
"D": {"start": 1.5, "min": 0.5, "max": 2.0, "step": 0.25}}
|
|
|
|
def test_add_param(self):
|
|
ms = project.ModelSpace()
|
|
ms.start['A'] = 2.1
|
|
ms.min['A'] = 2.0
|
|
ms.max['A'] = 3.0
|
|
ms.step['A'] = 0.05
|
|
ms.add_param("E", 5.0, 1.0, 9.0, 0.2)
|
|
ms.add_param("F", 8.0, width=6.0, step=0.5)
|
|
d_start = {'A': 2.1, 'E': 5.0, 'F': 8.0}
|
|
d_min = {'A': 2.0, 'E': 1.0, 'F': 5.0}
|
|
d_max = {'A': 3.0, 'E': 9.0, 'F': 11.0}
|
|
d_step = {'A': 0.05, 'E': 0.2, 'F': 0.5}
|
|
self.assertDictEqual(ms.start, d_start)
|
|
self.assertDictEqual(ms.min, d_min)
|
|
self.assertDictEqual(ms.max, d_max)
|
|
self.assertDictEqual(ms.step, d_step)
|
|
|
|
def test_get_param(self):
|
|
ms = project.ModelSpace()
|
|
ms.add_param("A", **self.d1['A'])
|
|
ms.add_param("B", **self.d1['B'])
|
|
result = ms.get_param('B')
|
|
expected = {'start': 15.0, 'min': 0.0, 'max': 30.0, 'step': 1.0}
|
|
self.assertIsInstance(result, project.ParamSpace)
|
|
self.assertEqual(result.start, expected['start'])
|
|
self.assertEqual(result.min, expected['min'])
|
|
self.assertEqual(result.max, expected['max'])
|
|
self.assertEqual(result.step, expected['step'])
|
|
|
|
def test_set_param_dict(self):
|
|
ms = project.ModelSpace()
|
|
ms.set_param_dict(self.d1)
|
|
ms.set_param_dict(self.d2)
|
|
d_start = {'C': 22.0, 'D': 1.5}
|
|
d_min = {'C': 15.0, 'D': 0.5}
|
|
d_max = {'C': 25.0, 'D': 2.0}
|
|
d_step = {'C': 1.0, 'D': 0.25}
|
|
self.assertDictEqual(ms.start, d_start)
|
|
self.assertDictEqual(ms.min, d_min)
|
|
self.assertDictEqual(ms.max, d_max)
|
|
self.assertDictEqual(ms.step, d_step)
|
|
|
|
|
|
class TestScanCreator(unittest.TestCase):
|
|
"""
|
|
test case for @ref pmsco.project.ScanCreator class
|
|
|
|
"""
|
|
def test_load_1(self):
|
|
"""
|
|
test the load method, case 1
|
|
|
|
test for:
|
|
- correct array expansion of an ['e', 'a'] scan.
|
|
- correct file name expansion with place holders and pathlib.Path objects.
|
|
"""
|
|
sc = project.ScanCreator()
|
|
sc.filename = Path("{test_p}", "twoatom_energy_alpha.etpai")
|
|
sc.positions = {
|
|
"e": "np.arange(10, 400, 5)",
|
|
"t": "0",
|
|
"p": "0",
|
|
"a": "np.linspace(-30, 30, 31)"
|
|
}
|
|
sc.emitter = "Cu"
|
|
sc.initial_state = "2p3/2"
|
|
|
|
p = Path(__file__).parent / ".." / "projects" / "twoatom"
|
|
dirs = {"test_p": p,
|
|
"test_s": str(p)}
|
|
|
|
result = sc.load(dirs=dirs)
|
|
|
|
self.assertEqual(result.mode, ['e', 'a'])
|
|
self.assertEqual(result.emitter, sc.emitter)
|
|
self.assertEqual(result.initial_state, sc.initial_state)
|
|
|
|
e = np.arange(10, 400, 5)
|
|
a = np.linspace(-30, 30, 31)
|
|
t = p = np.asarray([0])
|
|
np.testing.assert_array_equal(result.energies, e)
|
|
np.testing.assert_array_equal(result.thetas, t)
|
|
np.testing.assert_array_equal(result.phis, p)
|
|
np.testing.assert_array_equal(result.alphas, a)
|
|
|
|
self.assertTrue(Path(result.filename).is_file(), msg=f"file {result.filename} not found")
|
|
|
|
|
|
class TestScan(unittest.TestCase):
|
|
"""
|
|
test case for @ref pmsco.project.Scan class
|
|
|
|
"""
|
|
def test_import_scan_file(self):
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|
test_file = os.path.join(base_dir, "..", "projects", "twoatom", "twoatom_energy_alpha.etpai")
|
|
|
|
scan = project.Scan()
|
|
scan.import_scan_file(test_file, "C", "1s")
|
|
|
|
mode = ['e', 'a']
|
|
self.assertEqual(scan.mode, mode)
|
|
|
|
ae = np.arange(10, 1005, 5)
|
|
at = np.asarray([0])
|
|
ap = np.asarray([0])
|
|
aa = np.arange(-90, 91, 1)
|
|
|
|
np.testing.assert_array_almost_equal(scan.energies, ae)
|
|
np.testing.assert_array_almost_equal(scan.thetas, at)
|
|
np.testing.assert_array_almost_equal(scan.phis, ap)
|
|
np.testing.assert_array_almost_equal(scan.alphas, aa)
|
|
|
|
def test_define_scan(self):
|
|
scan = project.Scan()
|
|
p0 = np.asarray([20])
|
|
p1 = np.linspace(1, 4, 4)
|
|
p2 = np.linspace(11, 13, 3)
|
|
d = {'t': p1, 'e': p0, 'p': p2}
|
|
scan.define_scan(d, "C", "1s")
|
|
|
|
ae = np.asarray([20])
|
|
at = np.asarray([1, 2, 3, 4])
|
|
ap = np.asarray([11, 12, 13])
|
|
aa = np.asarray([0])
|
|
|
|
np.testing.assert_array_almost_equal(scan.energies, ae)
|
|
np.testing.assert_array_almost_equal(scan.thetas, at)
|
|
np.testing.assert_array_almost_equal(scan.phis, ap)
|
|
np.testing.assert_array_almost_equal(scan.alphas, aa)
|
|
|
|
re = np.ones(12) * 20
|
|
rt = np.asarray([1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4])
|
|
rp = np.asarray([11, 12, 13, 11, 12, 13, 11, 12, 13, 11, 12, 13])
|
|
ra = np.ones(12) * 0
|
|
|
|
np.testing.assert_array_almost_equal(scan.raw_data['e'], re)
|
|
np.testing.assert_array_almost_equal(scan.raw_data['t'], rt)
|
|
np.testing.assert_array_almost_equal(scan.raw_data['p'], rp)
|
|
np.testing.assert_array_almost_equal(scan.raw_data['a'], ra)
|
|
|
|
|
|
class TestProject(unittest.TestCase):
|
|
def setUp(self):
|
|
# before each test method
|
|
self.project = project.Project()
|
|
|
|
def tearDown(self):
|
|
# after each test method
|
|
pass
|
|
|
|
@classmethod
|
|
def setup_class(cls):
|
|
# before any methods in this class
|
|
pass
|
|
|
|
@classmethod
|
|
def teardown_class(cls):
|
|
# teardown_class() after any methods in this class
|
|
pass
|
|
|
|
@mock.patch('pmsco.data.load_data')
|
|
@mock.patch('pmsco.data.save_data')
|
|
def test_combine_domains(self, save_data_mock, load_data_mock):
|
|
self.project.scans.append(project.Scan())
|
|
|
|
parent_task = dispatch.CalculationTask()
|
|
parent_task.change_id(model=0, scan=0)
|
|
parent_task.model['wdom1'] = 0.5
|
|
|
|
child_tasks = [parent_task.copy()] * 2
|
|
for idx, task in enumerate(child_tasks):
|
|
task.change_id(domain=idx)
|
|
|
|
data1 = data.create_data(5, datatype='EI')
|
|
data1['e'] = np.arange(5)
|
|
data1['i'] = 10.
|
|
data2 = data1.copy()
|
|
data2['i'] = 10.
|
|
load_data_mock.side_effect = [data1, data2]
|
|
data3 = data1.copy()
|
|
data3['i'] = (10. + 0.5 * 10.) / 1.5
|
|
|
|
self.project.combine_domains(parent_task, child_tasks)
|
|
|
|
save_data_mock.assert_called()
|
|
args, kwargs = save_data_mock.call_args
|
|
result_data = args[1]
|
|
np.testing.assert_array_almost_equal(result_data['e'], data3['e'])
|
|
np.testing.assert_array_almost_equal(result_data['i'], data3['i'])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|