diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 3c6a523..b3a33d0 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -41,4 +41,4 @@ jobs: pip install . - name: Test with pytest run: | - pytest + python -m pytest tests diff --git a/qse/qbits.py b/qse/qbits.py index e223410..b968485 100644 --- a/qse/qbits.py +++ b/qse/qbits.py @@ -210,8 +210,6 @@ def __init__( celldisp = np.zeros(shape=(3, 1)) self.set_celldisp(celldisp) - # print(type(positions), len(positions), type(positions[0])) - # print(type(scaled_positions), len(scaled_positions), type(scaled_positions[0])) if positions is None: if scaled_positions is None: positions = np.zeros((len(self.arrays["labels"]), 3)) @@ -490,16 +488,20 @@ def set_array(self, name, a, dtype=None, shape=None): b[:] = a def has(self, name): - """Check for existence of array. + """ + Check for existence of array. name must be one of: 'tags', 'momenta', 'masses', 'initial_magmoms', - 'initial_charges'.""" + 'initial_charges'. + """ # XXX extend has to calculator properties return name in self.arrays def set_tags(self, tags): - """Set tags for all qbits. If only one tag is supplied, it is - applied to all qbits.""" + """ + Set tags for all qbits. If only one tag is supplied, it is + applied to all qbits. + """ if isinstance(tags, int): tags = [tags] * len(self) self.set_array("tags", tags, int, ()) @@ -512,8 +514,10 @@ def get_tags(self): return np.zeros(len(self), int) def set_positions(self, newpositions, apply_constraint=True): - """Set positions, honoring any constraints. To ignore constraints, - use *apply_constraint=False*.""" + """ + Set positions, honoring any constraints. To ignore constraints, + use *apply_constraint=False*. + """ if self.constraints and apply_constraint: newpositions = np.array(newpositions, float) for constraint in self.constraints: @@ -522,7 +526,8 @@ def set_positions(self, newpositions, apply_constraint=True): self.set_array("positions", newpositions, shape=(3,)) def get_positions(self, wrap=False, **wrap_kw): - """Get array of positions. + """ + Get array of positions. Parameters: @@ -955,17 +960,17 @@ def rotate(self, a, v, center=(0, 0, 0), rotate_cell=False): Parameters ---------- - a = None: + a : Angle that the qbits is rotated around the vector 'v'. 'a' can also be a vector and then 'a' is rotated into 'v'. - v: + v : Vector to rotate the qbits around. Vectors can be given as strings: 'x', '-x', 'y', ... . - center = (0, 0, 0): + center : The center is kept fixed under the rotation. Use 'COP' to fix the center of positions or 'COU' to fix the center of - cell. + cell. Defaults to = (0, 0, 0). rotate_cell = False: If true the cell is also rotated. @@ -985,10 +990,9 @@ def rotate(self, a, v, center=(0, 0, 0), rotate_cell=False): if not isinstance(a, numbers.Real): a, v = v, a - norm = np.linalg.norm v = string2vector(v) - normv = norm(v) + normv = np.linalg.norm(v) if normv == 0.0: raise ZeroDivisionError("Cannot rotate: norm(v) == 0") @@ -1004,19 +1008,19 @@ def rotate(self, a, v, center=(0, 0, 0), rotate_cell=False): normv2 = np.linalg.norm(v2) if normv2 == 0: raise ZeroDivisionError("Cannot rotate: norm(a) == 0") - v2 /= norm(v2) + v2 /= np.linalg.norm(v2) c = np.dot(v, v2) v = np.cross(v, v2) - s = norm(v) + s = np.linalg.norm(v) # In case *v* and *a* are parallel, np.cross(v, v2) vanish # and can't be used as a rotation axis. However, in this # case any rotation axis perpendicular to v2 will do. eps = 1e-7 if s < eps: v = np.cross((0, 0, 1), v2) - if norm(v) < eps: + if np.linalg.norm(v) < eps: v = np.cross((1, 0, 0), v2) - assert norm(v) >= eps + assert np.linalg.norm(v) >= eps elif s > 0: v /= s @@ -1048,23 +1052,23 @@ def _centering_as_array(self, center): return center def euler_rotate(self, phi=0.0, theta=0.0, psi=0.0, center=(0, 0, 0)): - """Rotate qbits via Euler angles (in degrees). + """ + Rotate qbits via Euler angles (in degrees). See e.g http://mathworld.wolfram.com/EulerAngles.html for explanation. - Parameters: - - center : - The point to rotate about. A sequence of length 3 with the - coordinates, or 'COP' to select center of positions or - 'COU' to select center of cell. - phi : + Parameters + ---------- + phi : float The 1st rotation angle around the z axis. - theta : + theta : float Rotation around the x axis. - psi : + psi : float 2nd rotation around the z axis. - + center : + The point to rotate about. A sequence of length 3 with the + coordinates, or 'COM' to select the center of mass, 'COP' to + select center of positions or 'COU' to select center of cell. """ center = self._centering_as_array(center) @@ -1072,9 +1076,7 @@ def euler_rotate(self, phi=0.0, theta=0.0, psi=0.0, center=(0, 0, 0)): theta *= pi / 180 psi *= pi / 180 - # First move the molecule to the origin In contrast to MATLAB, - # numpy broadcasts the smaller array to the larger row-wise, - # so there is no need to play with the Kronecker product. + # First move the molecule to the origin. rcoords = self.positions - center # First Euler rotation about z in matrix form D = np.array( @@ -1208,19 +1210,32 @@ def rotate_dihedral(self, a1, a2, a3, a4, angle=None, mask=None, indices=None): start = self.get_dihedral(a1, a2, a3, a4) self.set_dihedral(a1, a2, a3, a4, angle + start, mask, indices) - def get_angle(self, a1, a2, a3, mic=False): - """Get angle formed by three qbits. + def get_angle(self, index_1: int, index_2: int, index_3: int, mic: bool = False): + """ + Get the angle in degress formed by three qbits. - Calculate angle in degrees between the vectors a2->a1 and - a2->a3. + Parameters + ---------- + index_1 : int + The index of the first qubit. + index_2 : int + The index of the second qubit. + index_3 : int + The index of the third qubit. + mic : bool + Use mic=True to use the Minimum Image Convention and calculate the + angle across periodic boundaries. - Use mic=True to use the Minimum Image Convention and calculate the - angle across periodic boundaries. + Notes + ----- + Let x1, x2, x3 be the vectors describing the positions of the three + qubits. Then we calcule the angle between x1-x2 and x3-x2. """ - return self.get_angles([[a1, a2, a3]], mic=mic)[0] + return self.get_angles([[index_1, index_2, index_3]], mic=mic)[0] def get_angles(self, indices, mic=False): - """Get angle formed by three qbits for multiple groupings. + """ + Get angle formed by three qbits for multiple groupings. Calculate angle in degrees between vectors between qbits a2->a1 and a2->a3, where a1, a2, and a3 are in each row of indices. @@ -1238,19 +1253,16 @@ def get_angles(self, indices, mic=False): v12 = a1s - a2s v32 = a3s - a2s - cell = None - pbc = None - if mic: - cell = self.cell - pbc = self.pbc + return get_angles(v12, v32, cell=self.cell, pbc=self.pbc) - return get_angles(v12, v32, cell=cell, pbc=pbc) + return get_angles(v12, v32, cell=None, pbc=None) def set_angle( self, a1, a2=None, a3=None, angle=None, mask=None, indices=None, add=False ): - """Set angle (in degrees) formed by three qbits. + """ + Set angle (in degrees) formed by three qbits. Sets the angle between vectors *a2*->*a1* and *a2*->*a3*. @@ -1259,7 +1271,8 @@ def set_angle( Same usage as in :meth:`ase.Qbits.set_dihedral`. If *mask* and *indices* are given, *indices* overwrites *mask*. If *mask* and *indices* - are not set, only *a3* is moved.""" + are not set, only *a3* is moved. + """ if any(a is None for a in [a2, a3, angle]): raise ValueError("a2, a3, and angle must not be None") diff --git a/tests/qbits_test.py b/tests/qbits_test.py index 7a30bd0..1fe37e6 100644 --- a/tests/qbits_test.py +++ b/tests/qbits_test.py @@ -161,9 +161,159 @@ def test_del_item(indices): if isinstance(indices, int): indices = [indices] - assert isinstance(qbits, qse.Qbits) - assert len(qbits) == nqbits - len(indices) - assert np.allclose( - qbits.get_positions(), + _qbits_checker( + qbits, positions[[i for i in range(nqbits) if i not in indices]], + nqbits - len(indices), + ) + + +def test_rotate(): + """Simple checks for rotate.""" + unit_square = np.array( + [[1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]] + ) + qbits = qse.Qbits(positions=unit_square) + qbits.rotate(90, "z") + + assert not np.allclose(qbits.get_positions(), unit_square) + + unit_square_rt = np.array( + [[0.0, 1.0, 0.0], [-1.0, 1.0, 0.0], [-1.0, 0.0, 0.0], [0.0, 0.0, 0.0]] + ) + assert np.allclose(qbits.get_positions(), unit_square_rt) + + # Check that a centered square is invariant (with relabelling). + square_centered = np.array( + [[1.0, 1.0, 0.0], [1.0, -1.0, 0.0], [-1.0, 1.0, 0.0], [-1.0, -1.0, 0.0]] + ) + qbits = qse.Qbits(positions=square_centered) + qbits.rotate(90, "z") + assert np.allclose(qbits.get_positions(), square_centered[[2, 0, 3, 1]]) + + +@pytest.mark.parametrize("angle", [10, 20, 30]) +def test_rotate_square_z(angle): + """Test rotating a square system about the z axis.""" + positions = np.array( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 1.0]] + ) + qbits = qse.Qbits(positions=positions) + qbits.rotate(angle, "z") + + angle_rads = np.pi * angle / 180 + new_positions = np.array( + [ + [0.0, 0.0, 0.0], + [np.cos(angle_rads), np.sin(angle_rads), 0.0], + [0.0, 0.0, 1.0], + [np.cos(angle_rads), np.sin(angle_rads), 1.0], + ] + ) + assert np.allclose(qbits.get_positions(), new_positions) + + +@pytest.mark.parametrize("a", ["x", "-y", (0.0, 2.0, 3)]) +@pytest.mark.parametrize("v", ["z", (1.0, 1.0, 0.0)]) +@pytest.mark.parametrize("center", [(0.0, 0.0, 0.0), (-1, 3, 0.2)]) +def test_rotate_distances(a, v, center): + """Check a random rotation preserves distances.""" + positions = np.random.rand(4, 3) + qbits = qse.Qbits(positions=positions) + distances = qbits.get_all_distances() + + qbits.rotate(a, v, center) + + assert not np.allclose(qbits.get_positions(), positions) + assert np.allclose(qbits.get_all_distances(), distances) + + +def test_euler_rotate_and_rotate(): + """Test rotate and euler_rotate agree.""" + ### Note that they rotate in different directions (clockwise & anti.) + ### May want to fix this. + positions = np.random.rand(4, 3) + + qbits_1 = qse.Qbits(positions=positions) + qbits_1.rotate(34, "z") + + qbits_2 = qse.Qbits(positions=positions) + qbits_2.euler_rotate(-34) + + assert np.allclose(qbits_1.get_positions(), qbits_2.get_positions()) + + +@pytest.mark.parametrize("phi", [11.2, 45.0]) +@pytest.mark.parametrize("theta", [-3.0, 40]) +@pytest.mark.parametrize("psi", [18.2]) +@pytest.mark.parametrize("center", [(0.0, 0.0, 0.0), (-1, 3, 0.2)]) +def test_euler_rotate_distances(phi, theta, psi, center): + """Check a random euler rotation preserves distances.""" + positions = np.random.rand(4, 3) + qbits = qse.Qbits(positions=positions) + distances = qbits.get_all_distances() + + qbits.euler_rotate(phi, theta, psi, center) + + assert not np.allclose(qbits.get_positions(), positions) + assert np.allclose(qbits.get_all_distances(), distances) + + +@pytest.mark.parametrize("angle", [90, 36.0]) +def test_get_angle(angle): + """Test get_angle on a simple 3-qbit system.""" + angle_rads = np.pi * angle / 180 + positions = np.array( + [ + [np.cos(angle_rads), np.sin(angle_rads), 0.0], + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + ] ) + qbits = qse.Qbits(positions=positions) + assert np.isclose(qbits.get_angle(0, 1, 2), angle) + + # Angle should be invariant under global translation. + qbits.translate(-3.3) + assert np.isclose(qbits.get_angle(0, 1, 2), angle) + + # Angle should be invariant under global rotation. + qbits.euler_rotate(10, 20, -44.5) + assert np.isclose(qbits.get_angle(0, 1, 2), angle) + + +def test_get_angles(): + """Test get_angles.""" + indices = np.array( + [ + [0, 1, 2], + [0, 1, 3], + [0, 1, 4], + [0, 1, 2], + ] + ) + + qbits = qse.Qbits(positions=np.random.rand(6, 3)) + angles = qbits.get_angles(indices) + + assert angles.shape == (indices.shape[0],) + assert np.isclose(angles[0], angles[-1]) + + +@pytest.mark.parametrize("angle", [11.2, 36.0]) +@pytest.mark.parametrize("indices", [[0, 1, 2], [2, 4, 1], [0, 4, 2]]) +def test_set_angle(angle, indices): + """Test set_angle on a simple 5-qbit system.""" + positions = np.array( + [ + [0.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [1.0, 0.0, 0.0], + [1.0, 1.0, 0.0], + [0.5, 0.5, 1.0], + ] + ) + qbits = qse.Qbits(positions=positions) + assert not np.isclose(angle, qbits.get_angle(*indices)) + qbits.set_angle(*indices, angle) + assert np.isclose(angle, qbits.get_angle(*indices))