pmsco-public/tests/test_cluster.py

560 lines
22 KiB
Python

"""
@package tests.test_cluster
unit tests for pmsco.cluster
the purpose of these tests is to check whether the code runs as expected in a particular environment.
to run the tests, change to the directory which contains the tests directory, and execute =nosetests=.
@pre nose must be installed (python-nose package on Debian).
@author Matthias Muntwiler, matthias.muntwiler@psi.ch
@copyright (c) 2015-19 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
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from io import BytesIO
import math
import numpy as np
import unittest
import pmsco.cluster as mc
class TestClusterFunctions(unittest.TestCase):
def setUp(self):
# before each test method
pass
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
@staticmethod
def create_cube():
"""
create a cluster object with atoms on the corners, faces and body center of the unit cube.
the atom types are unique in an arbitrary sequence.
the emitter is at the origin, atom type 1.
@return: cluster.Cluster object.
"""
clu = mc.Cluster()
clu.add_atom(1, np.asarray([0, 0, 0]), 1)
clu.add_atom(2, np.asarray([1, 0, 0]), 0)
clu.add_atom(3, np.asarray([0, 1, 0]), 0)
clu.add_atom(4, np.asarray([0, 0, 1]), 0)
clu.add_atom(5, np.asarray([-1, 0, 0]), 0)
clu.add_atom(6, np.asarray([0, -1, 0]), 0)
clu.add_atom(7, np.asarray([0, 0, -1]), 0)
clu.add_atom(8, np.asarray([1, 1, 0]), 0)
clu.add_atom(9, np.asarray([0, 1, 1]), 0)
clu.add_atom(10, np.asarray([1, 0, 1]), 0)
clu.add_atom(11, np.asarray([-1, 1, 0]), 0)
clu.add_atom(12, np.asarray([0, -1, 1]), 0)
clu.add_atom(13, np.asarray([1, 0, -1]), 0)
clu.add_atom(14, np.asarray([-1, -1, 0]), 0)
clu.add_atom(15, np.asarray([0, -1, -1]), 0)
clu.add_atom(16, np.asarray([-1, 0, -1]), 0)
clu.add_atom(17, np.asarray([1, -1, 0]), 0)
clu.add_atom(18, np.asarray([0, 1, -1]), 0)
clu.add_atom(19, np.asarray([-1, 0, 1]), 0)
clu.add_atom(20, np.asarray([1, 1, 1]), 0)
clu.add_atom(21, np.asarray([-1, 1, 1]), 0)
clu.add_atom(22, np.asarray([1, -1, 1]), 0)
clu.add_atom(23, np.asarray([1, 1, -1]), 0)
clu.add_atom(24, np.asarray([-1, -1, -1]), 0)
clu.add_atom(25, np.asarray([1, -1, -1]), 0)
clu.add_atom(26, np.asarray([-1, 1, -1]), 0)
clu.add_atom(27, np.asarray([-1, -1, 1]), 0)
return clu
def test_numpy_extract(self):
"""
test array extraction code which should be compatible to numpy versions before and after 1.14.
numpy 1.14 introduces changes to multi-column indexing of structured arrays such as data[['x','y']].
first, it will return a view instead of a copy.
second, it will assign fields by position rather than by name.
the first change affects our cluster code in several places
where we extract XYZ coordinates from the cluster data array.
this test checks whether the new code works with a particular numpy version.
@return: None
"""
clu = self.create_cube()
xy3 = np.empty((clu.data.shape[0], 2), np.float32)
xy3[:, 0] = clu.data['x']
xy3[:, 1] = clu.data['y']
ctr = np.asarray((1.0, 0.0, 0.0))
dist = np.linalg.norm(xy3 - ctr[0:2], axis=1)
self.assertAlmostEqual(1.0, dist[0])
self.assertAlmostEqual(0.0, dist[1])
clu.clear()
xy3 = np.empty((clu.data.shape[0], 2), np.float32)
xy3[:, 0] = clu.data['x']
xy3[:, 1] = clu.data['y']
ctr = np.asarray((1.0, 0.0, 0.0))
dist = np.linalg.norm(xy3 - ctr[0:2], axis=1)
self.assertEqual(0, dist.shape[0])
def test_get_positions(self):
"""
check that we get an independent copy of the original data.
@return: None
"""
clu = self.create_cube()
pos = clu.get_positions()
self.assertEqual(clu.data.shape[0], pos.shape[0])
self.assertEqual(3, pos.shape[1])
self.assertEqual(np.float32, pos.dtype)
self.assertEqual(1.0, pos[1, 0])
self.assertEqual(0.0, pos[1, 1])
self.assertEqual(0.0, pos[1, 2])
pos[0, 0] = 15.0
self.assertEqual(0.0, clu.data['x'][0])
# empty cluster
clu.clear()
self.assertEqual(clu.data.shape[0], 0)
self.assertEqual(3, pos.shape[1])
self.assertEqual(np.float32, pos.dtype)
def test_set_positions(self):
clu = mc.Cluster()
clu.data = np.zeros(2, dtype=clu.dtype)
pos = np.array([[1., 2., 3.], [4., 5., 6.]])
clu.set_positions(pos)
self.assertEqual(1., clu.data['x'][0])
self.assertEqual(2., clu.data['y'][0])
self.assertEqual(3., clu.data['z'][0])
self.assertEqual(4., clu.data['x'][1])
self.assertEqual(5., clu.data['y'][1])
self.assertEqual(6., clu.data['z'][1])
def test_get_emitters(self):
clu = self.create_cube()
clu.set_emitter(idx=0)
clu.set_emitter(idx=9)
self.assertEqual(2, clu.get_emitter_count())
result = clu.get_emitters(['x', 'y', 'z', 't'])
expect = [(0., 0., 0., 1), (1., 0., 1., 10)]
self.assertEqual(expect, result)
def test_get_z_layers(self):
clu = mc.Cluster()
clu.add_atom(1, np.asarray([1, 0, 0.1]), 0)
clu.add_atom(2, np.asarray([0, 1, -0.3]), 0)
clu.add_atom(1, np.asarray([0, 1, -0.2]), 0)
clu.add_atom(1, np.asarray([1, 0, 0]), 1)
clu.add_atom(1, np.asarray([0, 1, -0.2001]), 0)
clu.add_atom(2, np.asarray([0, 1, -0.1]), 0)
clu.add_atom(1, np.asarray([0, 1, -0.1999]), 0)
layers = clu.get_z_layers(0.01)
np.testing.assert_allclose(layers, np.asarray([-0.3, -0.2, -0.1, 0.0, +0.1]), atol=0.001)
def test_get_center(self):
clu = mc.Cluster()
clu.add_atom(1, np.asarray([1, 0, 0]), 0)
clu.add_atom(2, np.asarray([0, 0, 1]), 0)
clu.add_atom(1, np.asarray([0, 1, 0]), 0)
clu.add_atom(2, np.asarray([-1, -1, -1]), 0)
v0 = np.asarray([0, 0, 0])
v1 = np.asarray([1/2, 1/2, 0])
v2 = np.asarray([-1/2, -1/2, 0])
v = clu.get_center()
np.testing.assert_allclose(v, v0, atol=0.001)
v = clu.get_center(element=1)
np.testing.assert_allclose(v, v1, atol=0.001)
v = clu.get_center(element="He")
np.testing.assert_allclose(v, v2, atol=0.001)
def test_relax(self):
clu = mc.Cluster()
clu.add_atom(1, np.asarray([1, 0, 1]), 0)
clu.add_atom(1, np.asarray([1, 0, 0]), 1)
clu.add_atom(2, np.asarray([0, 1, -1]), 0)
clu.add_atom(1, np.asarray([0, 1, -2]), 0)
clu.add_atom(2, np.asarray([0, 1, -3]), 0)
idx = clu.relax(-0.3, -0.1, 2)
np.testing.assert_almost_equal(idx, np.asarray([2, 4]))
np.testing.assert_allclose(clu.get_position(0), np.asarray([1, 0, 1]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(1), np.asarray([1, 0, 0]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(2), np.asarray([0, 1, -1.1]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(3), np.asarray([0, 1, -2.0]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(4), np.asarray([0, 1, -3.1]), atol=1e-6)
def test_rotate_x(self):
clu = mc.Cluster()
clu.add_atom(1, np.asarray([1, 0, 0]), 0)
clu.add_atom(1, np.asarray([0, 1, 0]), 0)
clu.add_atom(1, np.asarray([0, 0, 1]), 0)
clu.rotate_x(90)
np.testing.assert_allclose(clu.get_position(0), np.asarray([1, 0, 0]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(1), np.asarray([0, 0, 1]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(2), np.asarray([0, -1, 0]), atol=1e-6)
def test_rotate_y(self):
clu = mc.Cluster()
clu.add_atom(1, np.asarray([1, 0, 0]), 0)
clu.add_atom(1, np.asarray([0, 1, 0]), 0)
clu.add_atom(1, np.asarray([0, 0, 1]), 0)
clu.rotate_y(90)
np.testing.assert_allclose(clu.get_position(0), np.asarray([0, 0, -1]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(1), np.asarray([0, 1, 0]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(2), np.asarray([1, 0, 0]), atol=1e-6)
def test_rotate_z(self):
clu = mc.Cluster()
clu.add_atom(1, np.asarray([1, 0, 0]), 0)
clu.add_atom(1, np.asarray([0, 1, 0]), 0)
clu.add_atom(1, np.asarray([0, 0, 1]), 0)
clu.rotate_z(90)
np.testing.assert_allclose(clu.get_position(0), np.asarray([0, 1, 0]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(1), np.asarray([-1, 0, 0]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(2), np.asarray([0, 0, 1]), atol=1e-6)
def test_translate(self):
clu = mc.Cluster()
clu.add_atom(1, np.asarray([1, 0, 0]), 0)
clu.add_atom(2, np.asarray([0, 1, 0]), 0)
clu.add_atom(3, np.asarray([0, 0, 1]), 0)
v = np.array((0.1, 0.2, 0.3))
shift = clu.translate(v)
np.testing.assert_allclose(clu.get_position(0), np.asarray([1.1, 0.2, 0.3]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(1), np.asarray([0.1, 1.2, 0.3]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(2), np.asarray([0.1, 0.2, 1.3]), atol=1e-6)
np.testing.assert_allclose(shift, np.asarray([0, 1, 2]))
shift = clu.translate(v, element=3)
np.testing.assert_allclose(clu.get_position(0), np.asarray([1.1, 0.2, 0.3]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(1), np.asarray([0.1, 1.2, 0.3]), atol=1e-6)
np.testing.assert_allclose(clu.get_position(2), np.asarray([0.2, 0.4, 1.6]), atol=1e-6)
np.testing.assert_allclose(shift, np.asarray([2]))
def test_add_layer(self):
clu = mc.Cluster()
clu.set_rmax(2.0)
b1 = np.array((1.0, 0.0, 0.0))
b2 = np.array((0.0, 1.0, 0.0))
a1 = np.array((0.1, 0.0, -0.1))
clu.add_layer(7, a1, b1, b2)
exp_pos = [[-0.9, -1.0, -0.1], [0.1, -1.0, -0.1], [1.1, -1.0, -0.1],
[-1.9, 0.0, -0.1], [-0.9, 0.0, -0.1], [0.1, 0.0, -0.1], [1.1, 0.0, -0.1],
[-0.9, 1.0, -0.1], [0.1, 1.0, -0.1], [1.1, 1.0, -0.1]]
nn = len(exp_pos)
self.assertEqual(nn, clu.data.shape[0])
self.assertEqual(nn, clu.get_atom_count())
act_pos = np.sort(clu.get_positions(), axis=0)
exp_pos = np.sort(np.array(exp_pos), axis=0)
np.testing.assert_allclose(act_pos, exp_pos)
act_idx = np.unique(clu.data['i'])
self.assertEqual(nn, act_idx.shape[0])
act_typ = (clu.data['t'] == 7).nonzero()[0]
self.assertEqual(nn, act_typ.shape[0])
def test_add_bulk(self):
clu = mc.Cluster()
clu.set_rmax(2.0)
b1 = np.array((1.0, 0.0, 0.0))
b2 = np.array((0.0, 1.0, 0.0))
b3 = np.array((0.0, 0.0, 1.0))
a1 = np.array((0.1, 0.0, -0.1))
z_surf = 0.8
clu.add_bulk(7, a1, b1, b2, b3, z_surf=z_surf)
r_great = max(clu.rmax, np.linalg.norm(a1))
n1 = max(int(r_great / np.linalg.norm(b1)) + 1, 4) * 3
n2 = max(int(r_great / np.linalg.norm(b2)) + 1, 4) * 3
n3 = max(int(r_great / np.linalg.norm(b3)) + 1, 4) * 3
exp_pos = []
nn = 0
for i1 in range(-n1, n1 + 1):
for i2 in range(-n2, n2 + 1):
for i3 in range(-n3, n3 + 1):
v = a1 + b1 * i1 + b2 * i2 + b3 * i3
if np.linalg.norm(v) <= clu.rmax and v[2] <= z_surf:
exp_pos.append(v)
nn += 1
self.assertEqual(nn, clu.data.shape[0])
self.assertEqual(nn, clu.get_atom_count())
act_pos = np.sort(clu.get_positions(), axis=0)
exp_pos = np.sort(np.array(exp_pos), axis=0)
np.testing.assert_allclose(act_pos, exp_pos)
act_idx = np.unique(clu.data['i'])
self.assertEqual(nn, act_idx.shape[0])
act_typ = (clu.data['t'] == 7).nonzero()[0]
self.assertEqual(nn, act_typ.shape[0])
def test_add_cluster(self):
clu1 = mc.Cluster()
clu1.add_atom(1, np.asarray([0, 0, 0]), 1)
clu1.add_atom(2, np.asarray([1, 0, 0]), 0)
clu1.add_atom(3, np.asarray([0, 1, 0]), 0)
clu1.add_atom(4, np.asarray([0, 0, -1]), 0)
clu1.add_atom(5, np.asarray([0, 0, -2]), 0)
clu2 = mc.Cluster()
clu2.add_atom(3, np.asarray([-0.2, 0, 0]), 0) # unique
clu2.add_atom(4, np.asarray([0, -0.2, 0]), 0) # unique
clu2.add_atom(5, np.asarray([0, 0.05, -1]), 0) # not unique
clu2.add_atom(5, np.asarray([0, 0, -1.09]), 0) # just within tolerance of uniqueness
clu2.add_atom(6, np.asarray([0, 0, -1.99]), 0) # not unique
clu2.add_atom(7, np.asarray([0, 0, -1.10]), 0) # just out of tolerance of uniqueness
clu1.set_rmax(1.5)
clu1.add_cluster(clu2, check_rmax=True, check_unique=True, tol=0.1)
self.assertEqual(5+3, clu1.get_atom_count())
self.assertEqual(7, clu1.data['t'][-1])
self.assertEqual(6, clu2.data.shape[0])
def test_find_positions(self):
clu = mc.Cluster()
# from hbncu project
b_surf = 2.50
clu.set_rmax(b_surf * 10.0)
b1 = np.array((b_surf, 0.0, 0.0))
b2 = np.array((b_surf / 2.0, b_surf * math.sqrt(3.0) / 2.0, 0.0))
a_N = np.array((0.0, 0.0, 0.0))
a_B = np.array(((b1[0] + b2[0]) / 3.0, (b1[1] + b2[1]) / 3.0, 0.0))
emitter = a_N + b1 * 1 + b2 * 2
clu.add_layer(7, a_N, b1, b2)
clu.add_layer(5, a_B, b1, b2)
pos = clu.find_positions(pos=emitter)
self.assertEqual(1, len(pos))
self.assertEqual(206, pos[0])
# position in the format returned by get_emitters
emitter = (emitter[0], emitter[1], emitter[2], 7)
pos = clu.find_positions(pos=emitter)
self.assertEqual(1, len(pos))
self.assertEqual(206, pos[0])
def test_find_index_cylinder(self):
clu = self.create_cube()
pos = np.array((0.8, 0.8, 0.8))
rxy = 0.5
rz = 1.0
idx = clu.find_index_cylinder(pos, rxy, rz, None)
self.assertEqual(2, len(idx))
self.assertEqual(8, clu.get_atomtype(idx[0]))
self.assertEqual(20, clu.get_atomtype(idx[1]))
idx = clu.find_index_cylinder(pos, rxy, rz, 8)
self.assertEqual(1, len(idx))
def test_trim_cylinder(self):
clu = mc.Cluster()
clu.set_rmax(10.0)
v_pos = np.asarray([0, 0, 0])
v_lat1 = np.asarray([1, 0, 0])
v_lat2 = np.asarray([0, 1, 0])
v_lat3 = np.asarray([0, 0, 1])
clu.add_bulk(7, v_pos, v_lat1, v_lat2, v_lat3)
clu.set_emitter(pos=v_pos)
r0 = 2.3
z0 = 4.2
clu.trim_cylinder(r0, z0)
self.assertEqual(clu.dtype, clu.data.dtype)
self.assertEqual(21 * 5, clu.data.shape[0])
self.assertEqual(2, clu.data[1]['i'])
self.assertEqual('N', clu.data[1]['s'])
self.assertEqual(7, clu.data[1]['t'])
self.assertEqual(1, clu.get_emitter_count())
n_low = np.sum(clu.data['z'] < -z0)
self.assertEqual(0, n_low)
n_high = np.sum(clu.data['z'] > z0)
self.assertEqual(0, n_high)
n_out = np.sum(clu.data['x']**2 + clu.data['y']**2 > r0**2)
self.assertEqual(0, n_out)
def test_trim_sphere(self):
clu = mc.Cluster()
clu.set_rmax(10.0)
v_pos = np.asarray([0, 0, 0])
v_lat1 = np.asarray([1, 0, 0])
v_lat2 = np.asarray([0, 1, 0])
v_lat3 = np.asarray([0, 0, 1])
clu.add_bulk(7, v_pos, v_lat1, v_lat2, v_lat3)
clu.set_emitter(pos=v_pos)
r0 = 2.3
clu.trim_sphere(r0)
self.assertEqual(clu.dtype, clu.data.dtype)
self.assertEqual(39, clu.data.shape[0])
self.assertEqual(2, clu.data[1]['i'])
self.assertEqual('N', clu.data[1]['s'])
self.assertEqual(7, clu.data[1]['t'])
self.assertEqual(1, clu.get_emitter_count())
n_out = np.sum(clu.data['x']**2 + clu.data['y']**2 + clu.data['z'] > r0**2)
self.assertEqual(0, n_out)
def test_trim_paraboloid(self):
clu = mc.Cluster()
clu.set_rmax(10.0)
v_pos = np.asarray([0, 0, 0])
v_lat1 = np.asarray([1, 0, 0])
v_lat2 = np.asarray([0, 1, 0])
v_lat3 = np.asarray([0, 0, 1])
clu.add_bulk(7, v_pos, v_lat1, v_lat2, v_lat3)
clu.set_emitter(pos=v_pos)
r0 = 3.5
z0 = -2.3
clu.trim_paraboloid(r0, z0)
self.assertEqual(clu.data.dtype, clu.dtype)
self.assertEqual(63, clu.data.shape[0])
self.assertEqual(2, clu.data[1]['i'])
self.assertEqual('N', clu.data[1]['s'])
self.assertEqual(7, clu.data[1]['t'])
self.assertEqual(1, clu.get_emitter_count())
n_low = np.sum(clu.data['z'] < z0)
self.assertEqual(0, n_low)
n_out = np.sum(clu.data['x']**2 + clu.data['y']**2 > r0**2)
self.assertEqual(0, n_out)
def test_trim_slab(self):
clu = self.create_cube()
clu.trim_slab('x', 0.5, 1.1)
self.assertEqual(clu.dtype, clu.data.dtype)
self.assertEqual(9 * 2, clu.data.shape[0])
self.assertEqual(1, clu.get_emitter_count())
def test_save_to_file(self):
clu = self.create_cube()
f = BytesIO()
pos = np.asarray((-1, -1, 0))
clu.set_emitter(pos=pos)
clu.save_to_file(f, mc.FMT_PMSCO, "qwerty", emitters_only=True)
f.seek(0)
line = f.readline()
self.assertEqual(line, b"# index element symbol class x y z emitter charge\n", b"line 1: " + line)
line = f.readline()
self.assertRegexpMatches(line, b"[0-9]+ +1 +H +[0-9]+ +[0.]+ +[0.]+ +[0.]+ +1 +[0.]", b"line 3: " + line)
line = f.readline()
self.assertRegexpMatches(line, b"[0-9]+ +14 +Si +[0-9]+ +[01.-]+ +[01.-]+ +[0.]+ +1 +[0.]", b"line 4: " + line)
line = f.readline()
self.assertEqual(b"", line, b"end of file")
f = BytesIO()
clu.save_to_file(f, mc.FMT_XYZ, "qwerty", emitters_only=True)
f.seek(0)
line = f.readline()
self.assertEqual(b"2\n", line, b"line 1: " + line)
line = f.readline()
self.assertEqual(b"qwerty\n", line, b"line 2: " + line)
line = f.readline()
self.assertRegexpMatches(line, b"H +[0.]+ +[0.]+ +[0.]+", b"line 3: " + line)
line = f.readline()
self.assertRegexpMatches(line, b"Si +[01.-]+ +[01.-]+ +[0.]+", b"line 4: " + line)
line = f.readline()
self.assertEqual(b"", line, b"end of file")
def test_load_from_file(self):
f = BytesIO()
f.write(b"2\n")
f.write(b"qwerty\n")
f.write(b"H 0.5 0.6 0.7\n")
f.write(b"Si -1.5 -1.6 -1.7\n")
f.seek(0)
clu = mc.Cluster()
clu.load_from_file(f, fmt=mc.FMT_XYZ)
np.testing.assert_allclose(clu.data['t'], np.array([1, 14]))
np.testing.assert_allclose(clu.data['x'], np.array([0.5, -1.5]))
np.testing.assert_allclose(clu.data['y'], np.array([0.6, -1.6]))
np.testing.assert_allclose(clu.data['z'], np.array([0.7, -1.7]))
np.testing.assert_allclose(clu.data['e'], np.array([0, 0]))
np.testing.assert_allclose(clu.data['q'], np.array([0, 0]))
f = BytesIO()
f.write(b"# index element symbol class x y z emitter charge\n")
# ['i', 't', 's', 'c', 'x', 'y', 'z', 'e', 'q']
f.write(b"1 6 C 1 0.5 0.6 0.7 0 -0.5\n")
f.write(b"2 14 Si 2 -1.5 -1.6 -1.7 1 0.5\n")
f.seek(0)
clu = mc.Cluster()
clu.load_from_file(f, fmt=mc.FMT_PMSCO)
np.testing.assert_allclose(clu.data['t'], np.array([6, 14]))
np.testing.assert_allclose(clu.data['x'], np.array([0.5, -1.5]))
np.testing.assert_allclose(clu.data['y'], np.array([0.6, -1.6]))
np.testing.assert_allclose(clu.data['z'], np.array([0.7, -1.7]))
np.testing.assert_allclose(clu.data['e'], np.array([0, 1]))
np.testing.assert_allclose(clu.data['q'], np.array([-0.5, 0.5]))
def test_update_atoms(self):
clu = mc.Cluster()
clu.add_atom(1, np.asarray([0, 0, 0]), 1)
clu.add_atom(3, np.asarray([0, 1, 0]), 0)
clu.add_atom(5, np.asarray([-1, 0, 0]), 0)
clu.add_atom(6, np.asarray([0, -1, 0]), 0)
clu.add_atom(2, np.asarray([1, 0, 0]), 0)
clu.add_atom(4, np.asarray([0, 0, 1]), 0)
other = mc.Cluster()
other.add_atom(1, np.asarray([0, 0, 0]), 1)
other.add_atom(5, np.asarray([-1, 0, 0]), 0)
other.add_atom(2, np.asarray([1, 0, 0]), 0)
other.add_atom(6, np.asarray([0, -1, 0]), 0)
other.add_atom(3, np.asarray([0, 1, 0]), 0)
other.add_atom(4, np.asarray([0, 0, 1]), 0)
other.data['c'] = np.asarray((1, 2, 2, 3, 3, 4))
clu.update_atoms(other, {'c'})
expected = np.asarray((1, 3, 2, 3, 2, 4))
np.testing.assert_array_equal(expected, clu.data['c'])
def test_calc_scattering_angles(self):
clu = mc.Cluster()
ref_em = np.asarray([0.1, -0.1, 0.5])
ref_th = np.asarray([0., 15., 90., 100., 120.])
ref_ph = np.asarray([0., 90., 180., 270., 360.])
ref_di = np.asarray([0.5, 1.0, 1.5, 2.0, 2.5])
exp_th = ref_th[0:4]
exp_ph = ref_ph[0:4]
exp_di = ref_di[0:4]
sel_ph = exp_ph > 180.
exp_ph[sel_ph] = exp_ph[sel_ph] - 360.
idx_em = clu.add_atom(1, ref_em, 1)
for i, r in enumerate(ref_di):
v = np.asarray([
r * math.cos(math.radians(ref_ph[i])) * math.sin(math.radians(ref_th[i])) + ref_em[0],
r * math.sin(math.radians(ref_ph[i])) * math.sin(math.radians(ref_th[i])) + ref_em[1],
r * math.cos(math.radians(ref_th[i])) + ref_em[2]])
clu.add_atom(i, v, 0)
result = clu.calc_scattering_angles(idx_em, 2.2)
np.testing.assert_allclose(result['index'], np.arange(1, exp_di.shape[0] + 1))
np.testing.assert_allclose(result['polar'], exp_th, atol=1e-3)
np.testing.assert_allclose(result['azimuth'], exp_ph, atol=1e-3)
np.testing.assert_allclose(result['dist'], exp_di, rtol=1e-5)