1# -*- coding: utf-8 -*-
2# MolMod is a collection of molecular modelling tools for python.
3# Copyright (C) 2007 - 2019 Toon Verstraelen <Toon.Verstraelen@UGent.be>, Center
4# for Molecular Modeling (CMM), Ghent University, Ghent, Belgium; all rights
5# reserved unless otherwise stated.
6#
7# This file is part of MolMod.
8#
9# MolMod is free software; you can redistribute it and/or
10# modify it under the terms of the GNU General Public License
11# as published by the Free Software Foundation; either version 3
12# of the License, or (at your option) any later version.
13#
14# MolMod is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program; if not, see <http://www.gnu.org/licenses/>
21#
22# --
23"""Data structure & tools to work with periodic systems"""
24
25
26from __future__ import division
27
28from builtins import range
29import numpy as np
30
31from molmod.utils import cached, ReadOnly, ReadOnlyAttribute
32
33
34__all__ = ["UnitCell"]
35
36
37class UnitCell(ReadOnly):
38    """Extensible representation of a unit cell.
39
40       Most attributes of the UnitCell object are treated as constants. If you
41       want to modify the unit cell, just create a modified UnitCell
42       object. This facilitates the caching of derived quantities such as the
43       distance matrices, while it imposes a cleaner coding style without
44       a significant computational overhead.
45    """
46    eps = 1e-6 # small positive number, below this value is approximately zero
47    matrix = ReadOnlyAttribute(np.ndarray, none=False, npdim=2,
48        npshape=(3,3), npdtype=np.floating, doc="matrix whose columns are the "
49        "primitive cell vectors")
50    active = ReadOnlyAttribute(np.ndarray, none=False, npdim=1, npshape=(3,),
51        npdtype=np.bool_, doc="the active cell vectors")
52
53    def __init__(self, matrix, active=None):
54        """
55           Argument:
56            | ``matrix``  --  the array with cell vectors. each column
57                              corresponds to a single cell vector.
58
59           Optional argument:
60            | ``active``  --  an array with three boolean values indicating
61                              which cell vectors are active. [default: all three
62                              True]
63        """
64        if active is None:
65            active = np.array([True, True, True])
66        self.matrix = matrix
67        self.active = active
68        # sanity checks for the unit cell
69        for col, name in enumerate(["a", "b", "c"]):
70            if self.active[col]:
71                norm = np.linalg.norm(matrix[:, col])
72                if norm < self.eps:
73                    raise ValueError("The length of ridge %s is (nearly) zero." % name)
74        if abs(self.volume) < self.eps:
75            raise ValueError("The ridges of the unit cell are (nearly) linearly dependent vectors.")
76
77    def __mul__(self, x):
78        return self.copy_with(matrix=self.matrix*x)
79
80    def __truediv__(self, x):
81        return self.copy_with(matrix=self.matrix/x)
82
83    @classmethod
84    def from_parameters3(cls, lengths, angles):
85        """Construct a 3D unit cell with the given parameters
86
87           The a vector is always parallel with the x-axis and they point in the
88           same direction. The b vector is always in the xy plane and points
89           towards the positive y-direction. The c vector points towards the
90           positive z-direction.
91        """
92        for length in lengths:
93            if length <= 0:
94                raise ValueError("The length parameters must be strictly positive.")
95        for angle in angles:
96            if angle <= 0 or angle >= np.pi:
97                raise ValueError("The angle parameters must lie in the range ]0 deg, 180 deg[.")
98        del length
99        del angle
100
101        matrix = np.zeros((3, 3), float)
102
103        # first cell vector along x-axis
104        matrix[0, 0] = lengths[0]
105
106        # second cell vector in x-y plane
107        matrix[0, 1] = np.cos(angles[2])*lengths[1]
108        matrix[1, 1] = np.sin(angles[2])*lengths[1]
109
110        # Finding the third cell vector is slightly more difficult. :-)
111        # It works like this:
112        # The dot products of a with c, b with c and c with c are known. the
113        # vector a has only an x component, b has no z component. This results
114        # in the following equations:
115        u_a = lengths[0]*lengths[2]*np.cos(angles[1])
116        u_b = lengths[1]*lengths[2]*np.cos(angles[0])
117        matrix[0, 2] = u_a/matrix[0, 0]
118        matrix[1, 2] = (u_b - matrix[0, 1]*matrix[0, 2])/matrix[1, 1]
119        u_c = lengths[2]**2 - matrix[0, 2]**2 - matrix[1, 2]**2
120        if u_c < 0:
121            raise ValueError("The given cell parameters do not correspond to a unit cell.")
122        matrix[2, 2] = np.sqrt(u_c)
123
124        active = np.ones(3, bool)
125        return cls(matrix, active)
126
127    @cached
128    def volume(self):
129        """The volume of the unit cell
130
131           The actual definition of the volume depends on the number of active
132           directions:
133
134           * num_active == 0  --  always -1
135           * num_active == 1  --  length of the cell vector
136           * num_active == 2  --  surface of the parallelogram
137           * num_active == 3  --  volume of the parallelepiped
138        """
139        active = self.active_inactive[0]
140        if len(active) == 0:
141            return -1
142        elif len(active) == 1:
143            return np.linalg.norm(self.matrix[:, active[0]])
144        elif len(active) == 2:
145            return np.linalg.norm(np.cross(self.matrix[:, active[0]], self.matrix[:, active[1]]))
146        elif len(active) == 3:
147            return abs(np.linalg.det(self.matrix))
148
149    @cached
150    def active_inactive(self):
151        """The indexes of the active and the inactive cell vectors"""
152        active_indices = []
153        inactive_indices = []
154        for index, active in enumerate(self.active):
155            if active:
156                active_indices.append(index)
157            else:
158                inactive_indices.append(index)
159        return active_indices, inactive_indices
160
161    @cached
162    def reciprocal(self):
163        """The reciprocal of the unit cell
164
165           In case of a three-dimensional periodic system, this is trivially the
166           transpose of the inverse of the cell matrix. This means that each
167           column of the matrix corresponds to a reciprocal cell vector. In case
168           of lower-dimensional periodicity, the inactive columns are zero, and
169           the active columns span the same sub space as the original cell
170           vectors.
171        """
172        U, S, Vt = np.linalg.svd(self.matrix*self.active)
173        Sinv = np.zeros(S.shape, float)
174        for i in range(3):
175            if abs(S[i]) < self.eps:
176                Sinv[i] = 0.0
177            else:
178                Sinv[i] = 1.0/S[i]
179        return np.dot(U*Sinv, Vt)*self.active
180
181    @cached
182    def parameters(self):
183        """The cell parameters (lengths and angles)"""
184        length_a = np.linalg.norm(self.matrix[:, 0])
185        length_b = np.linalg.norm(self.matrix[:, 1])
186        length_c = np.linalg.norm(self.matrix[:, 2])
187        alpha = np.arccos(np.dot(self.matrix[:, 1], self.matrix[:, 2]) / (length_b * length_c))
188        beta = np.arccos(np.dot(self.matrix[:, 2], self.matrix[:, 0]) / (length_c * length_a))
189        gamma = np.arccos(np.dot(self.matrix[:, 0], self.matrix[:, 1]) / (length_a * length_b))
190        return (
191            np.array([length_a, length_b, length_c], float),
192            np.array([alpha, beta, gamma], float)
193        )
194
195    @cached
196    def ordered(self):
197        """An equivalent unit cell with the active cell vectors coming first"""
198        active, inactive = self.active_inactive
199        order = active + inactive
200        return UnitCell(self.matrix[:,order], self.active[order])
201
202    @cached
203    def alignment_a(self):
204        """Computes the rotation matrix that aligns the unit cell with the
205           Cartesian axes, starting with cell vector a.
206
207           * a parallel to x
208           * b in xy-plane with b_y positive
209           * c with c_z positive
210        """
211        from molmod.transformations import Rotation
212        new_x = self.matrix[:, 0].copy()
213        new_x /= np.linalg.norm(new_x)
214        new_z = np.cross(new_x, self.matrix[:, 1])
215        new_z /= np.linalg.norm(new_z)
216        new_y = np.cross(new_z, new_x)
217        new_y /= np.linalg.norm(new_y)
218        return Rotation(np.array([new_x, new_y, new_z]))
219
220    @cached
221    def alignment_c(self):
222        """Computes the rotation matrix that aligns the unit cell with the
223           Cartesian axes, starting with cell vector c.
224
225           * c parallel to z
226           * b in zy-plane with b_y positive
227           * a with a_x positive
228        """
229        from molmod.transformations import Rotation
230        new_z = self.matrix[:, 2].copy()
231        new_z /= np.linalg.norm(new_z)
232        new_x = np.cross(self.matrix[:, 1], new_z)
233        new_x /= np.linalg.norm(new_x)
234        new_y = np.cross(new_z, new_x)
235        new_y /= np.linalg.norm(new_y)
236        return Rotation(np.array([new_x, new_y, new_z]))
237
238    @cached
239    def spacings(self):
240        """Computes the distances between neighboring crystal planes"""
241        result_invsq = (self.reciprocal**2).sum(axis=0)
242        result = np.zeros(3, float)
243        for i in range(3):
244            if result_invsq[i] > 0:
245                result[i] = result_invsq[i]**(-0.5)
246        return result
247
248    def to_fractional(self, cartesian):
249        """Convert Cartesian to fractional coordinates
250
251           Argument:
252            | ``cartesian``  --  Can be a numpy array with shape (3, ) or with shape
253                                 (N, 3).
254
255           The return value has the same shape as the argument. This function is
256           the inverse of to_cartesian.
257        """
258        return np.dot(cartesian, self.reciprocal)
259
260    def to_cartesian(self, fractional):
261        """Converts fractional to Cartesian coordinates
262
263           Argument:
264            | ``fractional``  --  Can be a numpy array with shape (3, ) or with shape
265                                  (N, 3).
266
267           The return value has the same shape as the argument. This function is
268           the inverse of to_fractional.
269        """
270        return np.dot(fractional, self.matrix.transpose())
271
272    def shortest_vector(self, delta):
273        """Compute the relative vector under periodic boundary conditions.
274
275           Argument:
276            | ``delta``  --  the relative vector between two points
277
278           The return value is not necessarily the shortest possible vector,
279           but instead is the vector with fractional coordinates in the range
280           [-0.5,0.5[. This is most of the times the shortest vector between
281           the two points, but not always. (See commented test.) It is always
282           the shortest vector for orthorombic cells.
283        """
284        fractional = self.to_fractional(delta)
285        fractional = np.floor(fractional + 0.5)
286        return delta - self.to_cartesian(fractional)
287
288    def add_cell_vector(self, vector):
289        """Returns a new unit cell with an additional cell vector"""
290        act = self.active_inactive[0]
291        if len(act) == 3:
292            raise ValueError("The unit cell already has three active cell vectors.")
293        matrix = np.zeros((3, 3), float)
294        active = np.zeros(3, bool)
295        if len(act) == 0:
296            # Add the new vector
297            matrix[:, 0] = vector
298            active[0] = True
299            return UnitCell(matrix, active)
300
301        a = self.matrix[:, act[0]]
302        matrix[:, 0] = a
303        active[0] = True
304        if len(act) == 1:
305            # Add the new vector
306            matrix[:, 1] = vector
307            active[1] = True
308            return UnitCell(matrix, active)
309
310        b = self.matrix[:, act[1]]
311        matrix[:, 1] = b
312        active[1] = True
313        if len(act) == 2:
314            # Add the new vector
315            matrix[:, 2] = vector
316            active[2] = True
317            return UnitCell(matrix, active)
318
319    def get_radius_ranges(self, radius, mic=False):
320        """Return ranges of indexes of the interacting neighboring unit cells
321
322           Interacting neighboring unit cells have at least one point in their
323           box volume that has a distance smaller or equal than radius to at
324           least one point in the central cell. This concept is of importance
325           when computing pair wise long-range interactions in periodic systems.
326
327           The mic (stands for minimum image convention) option can be used to
328           change the behavior of this routine such that only neighboring cells
329           are considered that have at least one point withing a distance below
330           `radius` from the center of the reference cell.
331        """
332        result = np.zeros(3, int)
333        for i in range(3):
334            if self.spacings[i] > 0:
335                if mic:
336                    result[i] = np.ceil(radius/self.spacings[i]-0.5)
337                else:
338                    result[i] = np.ceil(radius/self.spacings[i])
339        return result
340
341    def get_radius_indexes(self, radius, max_ranges=None):
342        """Return the indexes of the interacting neighboring unit cells
343
344           Interacting neighboring unit cells have at least one point in their
345           box volume that has a distance smaller or equal than radius to at
346           least one point in the central cell. This concept is of importance
347           when computing pair wise long-range interactions in periodic systems.
348
349           Argument:
350            | ``radius``  --  the radius of the interaction sphere
351
352           Optional argument:
353            | ``max_ranges``  --  numpy array with three elements: The maximum
354                                  ranges of indexes to consider. This is
355                                  practical when working with the minimum image
356                                  convention to reduce the generated bins to the
357                                  minimum image. (see binning.py) Use -1 to
358                                  avoid such limitations. The default is three
359                                  times -1.
360
361        """
362        if max_ranges is None:
363            max_ranges = np.array([-1, -1, -1])
364        ranges = self.get_radius_ranges(radius)*2+1
365        mask = (max_ranges != -1) & (max_ranges < ranges)
366        ranges[mask] = max_ranges[mask]
367        max_size = np.product(self.get_radius_ranges(radius)*2 + 1)
368        indexes = np.zeros((max_size, 3), int)
369
370        from molmod.ext import unit_cell_get_radius_indexes
371        reciprocal = self.reciprocal*self.active
372        matrix = self.matrix*self.active
373        size = unit_cell_get_radius_indexes(
374            matrix, reciprocal, radius, max_ranges, indexes
375        )
376        return indexes[:size]
377