1"""Contains functions for doing the inverse and forward normal mode transforms.
2
3Copyright (C) 2013, Joshua More and Michele Ceriotti
4
5This program is free software: you can redistribute it and/or modify
6it under the terms of the GNU General Public License as published by
7the Free Software Foundation, either version 3 of the License, or
8(at your option) any later version.
9
10This program is distributed in the hope that it will be useful,
11but WITHOUT ANY WARRANTY; without even the implied warranty of
12MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13GNU General Public License for more details.
14
15You should have received a copy of the GNU General Public License
16along with this program. If not, see <http.//www.gnu.org/licenses/>.
17
18
19Classes:
20   nm_trans: Uses matrix multiplication to do normal mode transformations.
21   nm_rescale: Uses matrix multiplication to do ring polymer contraction
22      or expansion.
23   nm_fft: Uses fast-Fourier transforms to do normal modes transformations.
24
25Functions:
26   mk_nm_matrix: Makes a matrix to transform between the normal mode and bead
27      representations.
28   mk_rs_matrix: Makes a matrix to transform between one number of beads and
29      another. Higher normal modes in the case of an expansion are set to zero.
30"""
31
32__all__ = ['nm_trans', 'nm_rescale', 'nm_fft']
33
34import numpy as np
35from ipi.utils.messages import verbosity, info
36
37def mk_nm_matrix(nbeads):
38   """Gets the matrix that transforms from the bead representation
39   to the normal mode representation.
40
41   If we return from this function a matrix C, then we transform between the
42   bead and normal mode representation using q_nm = C . q_b, q_b = C.T . q_nm
43
44   Args:
45      nbeads: The number of beads.
46   """
47
48   b2nm = np.zeros((nbeads,nbeads))
49   b2nm[0,:] = np.sqrt(1.0)
50   for j in range(nbeads):
51      for i in range(1, nbeads/2+1):
52         b2nm[i,j] = np.sqrt(2.0)*np.cos(2*np.pi*j*i/float(nbeads))
53      for i in range(nbeads/2+1, nbeads):
54         b2nm[i,j] = np.sqrt(2.0)*np.sin(2*np.pi*j*i/float(nbeads))
55   if (nbeads%2) == 0:
56      b2nm[nbeads/2,0:nbeads:2] = 1.0
57      b2nm[nbeads/2,1:nbeads:2] = -1.0
58   return b2nm/np.sqrt(nbeads)
59
60def mk_rs_matrix(nb1, nb2):
61   """Gets the matrix that transforms a path with nb1 beads into one with
62   nb2 beads.
63
64   If we return from this function a matrix T, then we transform between the
65   system with nb1 bead and the system of nb2 beads using q_2 = T . q_1
66
67   Args:
68      nb1: The initial number of beads.
69      nb2: The final number of beads.
70   """
71
72   if (nb1 == nb2):
73      return np.identity(nb1,float)
74   elif (nb1 > nb2):
75      b1_nm = mk_nm_matrix(nb1)
76      nm_b2 = mk_nm_matrix(nb2).T
77
78      #builds the "reduction" matrix that picks the normal modes we want to keep
79      b1_b2 = np.zeros((nb2, nb1), float)
80      b1_b2[0,0] = 1.0
81      for i in range(1, nb2/2+1):
82         b1_b2[i,i] = 1.0
83         b1_b2[nb2-i, nb1-i] = 1.0
84      if (nb2 % 2 == 0):
85         #if we are contracting down to an even number of beads, then we have to
86         #pick just one of the last degenerate modes to match onto the single
87         #stiffest mode in the new path
88         b1_b2[nb2/2, nb1-nb2/2] = 0.0
89
90      rs_b1_b2 = np.dot(nm_b2, np.dot(b1_b2, b1_nm))
91      return rs_b1_b2*np.sqrt(float(nb2)/float(nb1))
92   else:
93      return mk_rs_matrix(nb2, nb1).T*(float(nb2)/float(nb1))
94
95
96class nm_trans:
97   """Helper class to perform beads <--> normal modes transformation.
98
99   Attributes:
100      _b2nm: The matrix to transform between the bead and normal mode
101         representations.
102      _nm2b: The matrix to transform between the normal mode and bead
103         representations.
104   """
105
106   def __init__(self, nbeads):
107      """Initializes nm_trans.
108
109      Args:
110         nbeads: The number of beads.
111      """
112
113      self._b2nm = mk_nm_matrix(nbeads)
114      self._nm2b = self._b2nm.T
115
116   def b2nm(self, q):
117      """Transforms a matrix to the normal mode representation.
118
119      Args:
120         q: A matrix with nbeads rows, in the bead representation.
121      """
122
123      return np.dot(self._b2nm,q)
124
125   def nm2b(self, q):
126      """Transforms a matrix to the bead representation.
127
128      Args:
129         q: A matrix with nbeads rows, in the normal mode representation.
130      """
131
132      return np.dot(self._nm2b,q)
133
134
135class nm_rescale:
136   """Helper class to rescale a ring polymer between different number of beads.
137
138   Attributes:
139      _b1tob2: The matrix to transform between a ring polymer with 'nbeads1'
140         beads and another with 'nbeads2' beads.
141      _b2tob1: The matrix to transform between a ring polymer with 'nbeads2'
142         beads and another with 'nbeads1' beads.
143   """
144
145   def __init__(self, nbeads1, nbeads2):
146      """Initializes nm_rescale.
147
148      Args:
149         nbeads1: The initial number of beads.
150         nbeads2: The rescaled number of beads.
151      """
152
153      self._b1tob2 = mk_rs_matrix(nbeads1,nbeads2)
154      self._b2tob1 = self._b1tob2.T*(float(nbeads1)/float(nbeads2))
155
156   def b1tob2(self, q):
157      """Transforms a matrix from one value of beads to another.
158
159      Args:
160         q: A matrix with nbeads1 rows, in the bead representation.
161      """
162
163      return np.dot(self._b1tob2,q)
164
165   def b2tob1(self, q):
166      """Transforms a matrix from one value of beads to another.
167
168      Args:
169         q: A matrix with nbeads2 rows, in the bead representation.
170      """
171
172      return np.dot(self._b2tob1,q)
173
174
175
176class nm_fft:
177   """Helper class to perform beads <--> normal modes transformation
178      using Fast Fourier transforms.
179
180   Attributes:
181      fft: The fast-Fourier transform function to transform between the
182         bead and normal mode representations.
183      ifft: The inverse fast-Fourier transform function to transform
184         between the normal mode and bead representations.
185      qdummy: A matrix to hold a copy of the bead positions to transform
186         them to the normal mode representation.
187      qnmdummy: A matrix to hold a copy of the normal modes to transform
188         them to the bead representation.
189      nbeads: The number of beads.
190      natoms: The number of atoms.
191   """
192
193   def __init__(self, nbeads, natoms):
194      """Initializes nm_trans.
195
196      Args:
197         nbeads: The number of beads.
198         natoms: The number of atoms.
199      """
200
201      self.nbeads = nbeads
202      self.natoms = natoms
203      try:
204         import pyfftw
205         info("Import of PyFFTW successful", verbosity.medium)
206         self.qdummy = pyfftw.n_byte_align_empty((nbeads, 3*natoms), 16, 'float32')
207         self.qnmdummy = pyfftw.n_byte_align_empty((nbeads//2+1, 3*natoms), 16, 'complex64')
208         self.fft = pyfftw.FFTW(self.qdummy, self.qnmdummy, axes=(0,), direction='FFTW_FORWARD')
209         self.ifft = pyfftw.FFTW(self.qnmdummy, self.qdummy, axes=(0,), direction='FFTW_BACKWARD')
210      except ImportError: #Uses standard numpy fft library if nothing better
211                          #is available
212         info("Import of PyFFTW unsuccessful, using NumPy library instead", verbosity.medium)
213         self.qdummy = np.zeros((nbeads,3*natoms), dtype='float32')
214         self.qnmdummy = np.zeros((nbeads//2+1,3*natoms), dtype='complex64')
215         def dummy_fft(self):
216            self.qnmdummy = np.fft.rfft(self.qdummy, axis=0)
217         def dummy_ifft(self):
218            self.qdummy = np.fft.irfft(self.qnmdummy, n=self.nbeads, axis=0)
219         self.fft = lambda: dummy_fft(self)
220         self.ifft = lambda: dummy_ifft(self)
221
222   def b2nm(self, q):
223      """Transforms a matrix to the normal mode representation.
224
225      Args:
226         q: A matrix with nbeads rows and 3*natoms columns,
227            in the bead representation.
228      """
229
230      if self.nbeads == 1:
231         return q
232      self.qdummy[:] = q
233      self.fft()
234      if self.nbeads == 2:
235         return self.qnmdummy.real/np.sqrt(self.nbeads)
236
237      nmodes = self.nbeads/2
238
239      self.qnmdummy /= np.sqrt(self.nbeads)
240      qnm = np.zeros(q.shape)
241      qnm[0,:] = self.qnmdummy[0,:].real
242
243      if self.nbeads % 2 == 0:
244         self.qnmdummy[1:-1,:] *= np.sqrt(2)
245         (qnm[1:nmodes,:], qnm[self.nbeads:nmodes:-1,:]) = (self.qnmdummy[1:-1,:].real, self.qnmdummy[1:-1,:].imag)
246         qnm[nmodes,:] = self.qnmdummy[nmodes,:].real
247      else:
248         self.qnmdummy[1:,:] *= np.sqrt(2)
249         (qnm[1:nmodes+1,:], qnm[self.nbeads:nmodes:-1,:]) = (self.qnmdummy[1:,:].real, self.qnmdummy[1:,:].imag)
250
251      return qnm
252
253   def nm2b(self, qnm):
254      """Transforms a matrix to the bead representation.
255
256      Args:
257         qnm: A matrix with nbeads rows and 3*natoms columns,
258            in the normal mode representation.
259      """
260
261      if self.nbeads == 1:
262         return qnm
263      if self.nbeads == 2:
264         self.qnmdummy[:] = qnm
265         self.ifft()
266         return self.qdummy*np.sqrt(self.nbeads)
267
268      nmodes = self.nbeads/2
269      odd = self.nbeads - 2*nmodes  # 0 if even, 1 if odd
270
271      qnm_complex = np.zeros((nmodes+1, len(qnm[0,:])), complex)
272      qnm_complex[0,:] = qnm[0,:]
273      if not odd:
274         (qnm_complex[1:-1,:].real, qnm_complex[1:-1,:].imag) = (qnm[1:nmodes,:], qnm[self.nbeads:nmodes:-1,:])
275         qnm_complex[1:-1,:] /= np.sqrt(2)
276         qnm_complex[nmodes,:] = qnm[nmodes,:]
277      else:
278         (qnm_complex[1:,:].real, qnm_complex[1:,:].imag) = (qnm[1:nmodes+1,:], qnm[self.nbeads:nmodes:-1,:])
279         qnm_complex[1:,:] /= np.sqrt(2)
280
281      self.qnmdummy[:] = qnm_complex
282      self.ifft()
283      return self.qdummy*np.sqrt(self.nbeads)
284