1###########################################################################
2#                                                                         #
3#    physical_validation,                                                 #
4#    a python package to test the physical validity of MD results         #
5#                                                                         #
6#    Written by Michael R. Shirts <michael.shirts@colorado.edu>           #
7#               Pascal T. Merz <pascal.merz@colorado.edu>                 #
8#                                                                         #
9#    Copyright (C) 2012 University of Virginia                            #
10#              (C) 2017 University of Colorado Boulder                    #
11#                                                                         #
12#    This library is free software; you can redistribute it and/or        #
13#    modify it under the terms of the GNU Lesser General Public           #
14#    License as published by the Free Software Foundation; either         #
15#    version 2.1 of the License, or (at your option) any later version.   #
16#                                                                         #
17#    This library is distributed in the hope that it will be useful,      #
18#    but WITHOUT ANY WARRANTY; without even the implied warranty of       #
19#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU    #
20#    Lesser General Public License for more details.                      #
21#                                                                         #
22#    You should have received a copy of the GNU Lesser General Public     #
23#    License along with this library; if not, write to the                #
24#    Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,     #
25#    Boston, MA 02110-1301 USA                                            #
26#                                                                         #
27###########################################################################
28r"""
29This module contains low-level functionality of the
30`physical_validation.kinetic_energy` module. The functions in this module
31should generally not be called directly. Please use the high-level
32functions from `physical_validation.kinetic energy`.
33"""
34from __future__ import print_function
35from __future__ import division
36
37import scipy.stats as stats
38import numpy as np
39import multiprocessing as mproc
40
41from ..util import trajectory
42from . import plot
43
44
45def isclose(a, b, rel_tol=1e-09, abs_tol=0.0):
46    return abs(a-b) <= max(rel_tol * max(abs(a), abs(b)), abs_tol)
47
48
49def temperature(kin, ndof, kb=8.314e-3):
50    r"""
51    Calculates the temperature acccording to the equipartition theorem.
52
53    .. math::
54        T(K) = \frac{2K}{N k_B}
55
56    Parameters
57    ----------
58    kin : float
59        Kinetic energy.
60    ndof : float
61        Number of degrees of freedom.
62    kb : float
63        Boltzmann constant :math:`k_B`. Default: 8.314e-3 (kJ/mol).
64
65    Returns
66    -------
67    temperature : float
68        Calculated temperature.
69    """
70    # ndof * kb * T = 2 * kin
71    if isclose(ndof, 0):
72        return 0
73    return 2 * float(kin) / (float(ndof) * float(kb))
74
75
76def check_mb_ensemble(kin, temp, ndof, alpha, kb=8.314e-3, verbosity=1,
77                      screen=False, filename=None, ene_unit=None):
78    r"""
79    Checks if a kinetic energy trajectory is Maxwell-Boltzmann distributed.
80
81    .. warning: This is a low-level function. Additionally to being less
82       user-friendly, there is a higher probability of erroneous and / or
83       badly documented behavior due to unexpected inputs. Consider using
84       the high-level version based on the SimulationData object. See
85       physical_validation.kinetic_energy.check_mb_ensemble for more
86       information and full documentation.
87
88    Parameters
89    ----------
90    kin : array-like
91        Kinetic energy snapshots of the system.
92    temp : float
93        Target temperature of the system. Used to construct the
94        Maxwell-Boltzmann distribution.
95    ndof : float
96        Number of degrees of freedom in the system. Used to construct the
97        Maxwell-Boltzmann distribution.
98    alpha : float
99        Confidence. TODO: Check proper statistical definition.
100    kb : float
101        Boltzmann constant :math:`k_B`. Default: 8.314e-3 (kJ/mol).
102    verbosity : int
103        0: Silent.
104        1: Print result details.
105        2: Print additional information.
106        Default: False.
107    screen : bool
108        Plot distributions on screen. Default: False.
109    filename : string
110        Plot distributions to `filename`.pdf. Default: None.
111    ene_unit : string
112        Energy unit - used for plotting only.
113
114    Returns
115    -------
116    result : float
117        The p value of the test.
118
119    See Also
120    --------
121    physical_validation.kinetic_energy.check_mb_ensemble : High-level version
122    """
123
124    # Discard burn-in period and time-correlated frames
125    kin = trajectory.prepare(kin, verbosity=verbosity, name='Kinetic energy')
126
127    kt = kb * temp
128    d, p = stats.kstest(kin, 'chi2', (ndof, 0, kt/2))
129
130    do_plot = screen or filename is not None
131
132    if do_plot:
133        hist_sim, k_sim = np.histogram(kin, bins=25, normed=True)
134        k_sim = (k_sim[1:] + k_sim[:-1])/2
135
136        ana_dist = stats.chi2(df=ndof, scale=kt/2)
137        k_ana = np.linspace(ana_dist.ppf(0.0001),
138                            ana_dist.ppf(0.9999), 100)
139        hist_ana = ana_dist.pdf(k_ana)
140
141        data = [{'x': k_sim,
142                 'y': hist_sim * 100,
143                 'name': 'Simulation result'},
144                {'x': k_ana,
145                 'y': hist_ana * 100,
146                 'name': 'Maxwell-Boltzmann'}]
147
148        unit = ''
149        if ene_unit is not None:
150            unit = ' [' + ene_unit + ']'
151
152        plot.plot(data,
153                  legend='best',
154                  title='Simulation vs. Maxwell-Boltzmann',
155                  xlabel='Kinetic energy' + unit,
156                  ylabel='Probability [%]',
157                  sci_x=True,
158                  filename=filename,
159                  screen=screen)
160
161    if verbosity > 0:
162        message = ('Kolmogorov-Smirnov test result: p = {:g}\n'
163                   'Null hypothesis: Kinetic energy is Maxwell-Boltzmann distributed'.format(p))
164        if alpha is not None:
165            if p >= alpha:
166                message += ('\nConfidence alpha = {:f}\n'
167                            'Result: Hypothesis stands'.format(alpha))
168            elif p < alpha:
169                message += ('\nConfidence alpha = {:f}\n'
170                            'Result: Hypothesis rejected'.format(alpha))
171        print(message)
172
173    return p
174
175
176def check_equipartition(positions, velocities, masses,
177                        molec_idx, molec_nbonds,
178                        natoms, nmolecs,
179                        ndof_reduction_tra=0, ndof_reduction_rot=0,
180                        dtemp=0.1, temp=None, alpha=0.05,
181                        molec_groups=None,
182                        random_divisions=0, random_groups=2,
183                        ndof_molec=None, kin_molec=None,
184                        verbosity=2,
185                        screen=False, filename=None):
186    r"""
187    Checks the equipartition of a simulation trajectory.
188
189    .. warning: This is a low-level function. Additionally to being less
190       user-friendly, there is a higher probability of erroneous and / or
191       badly documented behavior due to unexpected inputs. Consider using
192       the high-level version based on the SimulationData object. See
193       physical_validation.kinetic_energy.check_equipartition for more
194       information and full documentation.
195
196    Parameters
197    ----------
198    positions : array-like (nframes x natoms x 3)
199        3d array containing the positions of all atoms for all frames
200    velocities : array-like (nframes x natoms x 3)
201        3d array containing the velocities of all atoms for all frames
202    masses : array-like (natoms x 1)
203        1d array containing the masses of all atoms
204    molec_idx : array-like (nmolecs x 1)
205        Index of first atom for every molecule
206    molec_nbonds : array-like (nmolecs x 1)
207        Number of bonds for every molecule
208    natoms : int
209        Number of atoms in the system
210    nmolecs : int
211        Number of molecules in the system
212    ndof_reduction_tra : int, optional
213        Number of center-of-mass translational degrees of freedom to
214        remove. Default: 0.
215    ndof_reduction_rot : int, optional
216        Number of center-of-mass rotational degrees of freedom to remove.
217        Default: 0.
218    dtemp : float, optional
219        Fraction of temperature deviation tolerated between groups.
220        Default: 0.05 (5%).
221    temp : float, optional
222        Target temperature of the simulation. If None, the kinetic
223        energies will not be tested for Maxwell-Boltzmann distribution,
224        but only compared amongst each others. Default: None.
225    alpha : float, optional
226        Confidence for Maxwell-Boltzmann test. Default: 0.05 (5%).
227    molec_groups : List[array-like] (ngroups x ?), optional
228        List of 1d arrays containing molecule indeces defining groups.
229        Useful to pre-define groups of molecules (e.g. solute / solvent,
230        liquid mixture species, ...). If None, no pre-defined molecule
231        groups will be tested. Default: None.
232
233        *Note:* If an empty 1d array is found as last element in the list, the remaining
234        molecules are collected in this array. This allows, for example, to only
235        specify the solute, and indicate the solvent by giving an empty array.
236    random_divisions : int, optional
237        Number of random division tests attempted. Default: 0 (random
238        division tests off).
239    random_groups : int, optional
240        Number of groups the system is randomly divided in. Default: 2.
241    ndof_molec : List[dict], optional
242        Pass in the degrees of freedom per molecule. Slightly increases speed of repeated
243        analysis of the same simulation run.
244    kin_molec : List[List[dict]], optional
245        Pass in the kinetic energy per molecule. Greatly increases speed of repeated
246        analysis of the same simulation run.
247    verbosity : int, optional
248        Verbosity level, where 0 is quiet and 3 very chatty. Default: 2.
249    screen : bool
250        Plot distributions on screen. Default: False.
251    filename : string
252        Plot distributions to `filename`.pdf. Default: None.
253
254    Returns
255    -------
256    result : int
257        Number of equipartition violations. Tune up verbosity for details.
258    ndof_molec : List[dict]
259        List of the degrees of freedom per molecule. Can be saved to increase speed of
260        repeated analysis of the same simulation run.
261    kin_molec : List[List[dict]]
262        List of the kinetic energy per molecule per frame. Can be saved to increase speed
263        of repeated analysis of the same simulation run.
264
265    See Also
266    --------
267    physical_validation.kinetic_energy.check_equipartition : High-level version
268
269    """
270
271    dict_keys = ['tot', 'tra', 'rni', 'rot', 'int']
272
273    # for each molecule, calculate total / translational / rotational & internal /
274    #   rotational / internal degrees of freedom
275    #   returns: List[dict] of floats (shape: nmolecs x 5 x 1)
276    if ndof_molec is None:
277        ndof_molec = calc_ndof(natoms, nmolecs, molec_idx, molec_nbonds,
278                               ndof_reduction_tra, ndof_reduction_rot)
279
280    # for each frame, calculate total / translational / rotational & internal /
281    #   rotational / internal kinetic energy for each molecule
282    if kin_molec is None:
283        try:
284            with mproc.Pool() as p:
285                kin_molec = p.starmap(calc_molec_kinetic_energy,
286                                      [(r, v, masses, molec_idx, natoms, nmolecs)
287                                       for r, v in zip(positions, velocities)])
288        except AttributeError:
289            # Parallel execution doesn't work in py2.7 for quite a number of reasons.
290            # Attribute error when opening the `with` region is the first error (and
291            # an easy one), but by far not the last. So let's just resort to non-parallel
292            # execution:
293            kin_molec = [calc_molec_kinetic_energy(r, v, masses, molec_idx, natoms, nmolecs)
294                         for r, v in zip(positions, velocities)]
295
296    result = []
297
298    # test system-wide tot, tra, rni, rot, int
299    if temp is not None:
300        # check for Maxwell-Boltzmann distribution of
301        # partitioned kinetic energy trajectories
302        result.extend(test_mb_dist(kin_molec=kin_molec,
303                                   ndof_molec=ndof_molec,
304                                   nmolecs=nmolecs,
305                                   temp=temp,
306                                   alpha=alpha,
307                                   dict_keys=dict_keys,
308                                   verbosity=verbosity,
309                                   screen=screen,
310                                   filename=filename))
311    else:
312        # compare partitioned temperatures to total temperature
313        result.extend(test_temp_diff(kin_molec=kin_molec,
314                                     ndof_molec=ndof_molec,
315                                     nmolecs=nmolecs,
316                                     dtemp=dtemp,
317                                     dict_keys=dict_keys,
318                                     verbosity=verbosity,
319                                     screen=screen,
320                                     filename=filename))
321    # divide in random groups
322    for i in range(random_divisions):
323        # randomly assign a group index to each molecule
324        group_idx = np.random.randint(random_groups, size=nmolecs)
325        # create molecule index for each group
326        groups = []
327        for rg in range(random_groups):
328            groups.append(np.arange(nmolecs)[group_idx == rg])
329        # test each group separately
330        for rg, group in enumerate(groups):
331            if verbosity > 0:
332                print('Testing randomly divided group {:d}'.format(rg))
333            if verbosity > 3:
334                print(group)
335            if temp is not None:
336                result.extend(test_mb_dist(kin_molec, ndof_molec, nmolecs, temp,
337                                           alpha, dict_keys, group, verbosity))
338            else:
339                result.extend(test_temp_diff(kin_molec, ndof_molec, nmolecs,
340                                             dtemp, dict_keys, group, verbosity))
341        # test groups against each others
342        for rg1, group1 in enumerate(groups):
343            for group2 in groups[rg1 + 1:]:
344                result.extend(test_temp_diff_groups(kin_molec, ndof_molec, nmolecs,
345                                                    group1, group2,
346                                                    dtemp, dict_keys, verbosity))
347
348    # use predefined group division?
349    # if no groups, return
350    if not molec_groups:
351        return result, ndof_molec, kin_molec
352    # is last group empty?
353    last_empty = not molec_groups[-1]
354    # get rid of empty groups
355    molec_groups = [g for g in molec_groups if g]
356    # if no groups now (all were empty), return now
357    if not molec_groups:
358        return result, ndof_molec, kin_molec
359
360    if last_empty:
361        # last group is [] -> insert remaining molecules
362        combined = []
363        for group in molec_groups:
364            combined.extend(group)
365        molec_groups[-1] = [m for m in range(nmolecs) if m not in combined]
366
367    for mg, group in enumerate(molec_groups):
368        if verbosity > 0:
369            print('Testing predifined divided group {:d}'.format(mg))
370        if verbosity > 3:
371            print(group)
372        if temp is not None:
373            result.extend(test_mb_dist(kin_molec, ndof_molec, nmolecs, temp,
374                                       alpha, dict_keys, group, verbosity))
375        else:
376            result.extend(test_temp_diff(kin_molec, ndof_molec, nmolecs,
377                                         dtemp, dict_keys, group, verbosity))
378    # test groups against each others
379    if len(molec_groups) > 1:
380        for rg1, group1 in enumerate(molec_groups):
381            for group2 in molec_groups[rg1 + 1:]:
382                result.extend(test_temp_diff_groups(kin_molec, ndof_molec, nmolecs,
383                                                    group1, group2,
384                                                    dtemp, dict_keys, verbosity))
385
386    return result, ndof_molec, kin_molec
387
388
389def calc_system_ndof(natoms, nmolecs, nbonds,
390                     stop_com_tra, stop_com_rot):
391    r"""
392    Calculates the total / translational / rotational & internal /
393    rotational / internal degrees of freedom of the system.
394
395    Parameters
396    ----------
397    natoms : int
398        Total number of atoms in the system
399    nmolecs : int
400        Total number of molecules in the system
401    nbonds : int
402        Total number of bonds in the system
403    stop_com_tra : bool
404        Was the center-of-mass translation removed during the simulation?
405    stop_com_rot : bool
406        Was the center-of-mass translation removed during the simulation?
407
408    Returns
409    -------
410    ndof : dict
411        Dictionary containing the degrees of freedom.
412        Keys: ['tot', 'tra', 'rni', 'rot', 'int']
413    """
414    # total ndof
415    ndof_tot = 3*natoms - nbonds
416
417    # ndof reduction due to COM motion constraining
418    if stop_com_tra:
419        ndof_tot -= 3
420    if stop_com_rot:
421        ndof_tot -= 3
422
423    # translational ndof
424    ndof_tot_tra = 3*nmolecs
425    if stop_com_tra:
426        ndof_tot -= 3
427
428    # rotational & internal ndof
429    ndof_tot_rni = ndof_tot - ndof_tot_tra
430
431    # rotational ndof
432    ndof_tot_rot = 3*nmolecs
433    if stop_com_tra:
434        ndof_tot -= 3
435
436    # internal ndof
437    ndof_tot_int = ndof_tot_rni - ndof_tot_rot
438
439    # return dict
440    ndof = {'tot': ndof_tot,
441            'tra': ndof_tot_tra,
442            'rni': ndof_tot_rni,
443            'rot': ndof_tot_rot,
444            'int': ndof_tot_int}
445    return ndof
446
447
448def calc_ndof(natoms, nmolecs,
449              molec_idx, molec_nbonds,
450              ndof_reduction_tra, ndof_reduction_rot):
451    r"""
452    Calculates the total / translational / rotational & internal /
453    rotational / internal degrees of freedom per molecule.
454
455    Parameters
456    ----------
457    natoms : int
458        Total number of atoms in the system
459    nmolecs : int
460        Total number of molecules in the system
461    molec_idx : List[int]
462        Index of first atom for every molecule
463    molec_nbonds : List[int]
464        Number of bonds for every molecule
465    ndof_reduction_tra : int
466        Number of center-of-mass translational degrees of freedom to
467        remove. Default: 0.
468    ndof_reduction_rot : int
469        Number of center-of-mass rotational degrees of freedom to remove.
470        Default: 0.
471
472    Returns
473    -------
474    ndof_molec : List[dict]
475        List of dictionaries containing the degrees of freedom for each molecule
476        Keys: ['tot', 'tra', 'rni', 'rot', 'int']
477    """
478    # check whether there are monoatomic molecules:
479    nmono = (np.array(molec_idx[1:]) - np.array(molec_idx[:-1]) == 1).sum()
480
481    # ndof to be deducted per molecule
482    # ndof reduction due to COM motion constraining
483    ndof_com_tra_pm = ndof_reduction_tra / nmolecs
484    ndof_com_rot_pm = ndof_reduction_rot / (nmolecs - nmono)
485
486    ndof_molec = []
487    # add last idx to molec_idx to ease looping
488    molec_idx = np.append(molec_idx, [natoms])
489    # loop over molecules
490    for idx_molec, (idx_atm_init, idx_atm_end) in enumerate(zip(molec_idx[:-1], molec_idx[1:])):
491        natoms = idx_atm_end - idx_atm_init
492        nbonds = molec_nbonds[idx_molec]
493        ndof_tot = 3*natoms - nbonds - ndof_com_tra_pm - ndof_com_rot_pm
494        ndof_tra = 3 - ndof_com_tra_pm
495        ndof_rni = ndof_tot - ndof_tra
496        ndof_rot = 3 - ndof_com_rot_pm
497        ndof_int = ndof_tot - ndof_tra - ndof_rot
498        if isclose(ndof_int, 0, abs_tol=1e-09):
499            ndof_int = 0
500        if natoms == 1:
501            ndof_tot = 3 - ndof_com_tra_pm
502            ndof_tra = 3 - ndof_com_tra_pm
503            ndof_rni = 0
504            ndof_rot = 0
505            ndof_int = 0
506        ndof_molec.append({'tot': ndof_tot,
507                           'tra': ndof_tra,
508                           'rni': ndof_rni,
509                           'rot': ndof_rot,
510                           'int': ndof_int})
511
512    return ndof_molec
513
514
515def calc_molec_kinetic_energy(pos, vel, masses,
516                              molec_idx, natoms, nmolecs):
517    r"""
518    Calculates the total / translational / rotational & internal /
519    rotational / internal kinetic energy per molecule.
520
521    Parameters
522    ----------
523    pos : nd-array (natoms x 3)
524        2d array containing the positions of all atoms
525    vel : nd-array (natoms x 3)
526        2d array containing the velocities of all atoms
527    masses : nd-array (natoms x 1)
528        1d array containing the masses of all atoms
529    molec_idx : nd-array (nmolecs x 1)
530        Index of first atom for every molecule
531    natoms : int
532        Total number of atoms in the system
533    nmolecs : int
534        Total number of molecules in the system
535
536    Returns
537    -------
538    kin : List[dict]
539        List of dictionaries containing the kinetic energies for each molecule
540        Keys: ['tot', 'tra', 'rni', 'rot', 'int']
541    """
542    # add last idx to molec_idx to ease looping
543    molec_idx = np.append(molec_idx, [natoms])
544
545    # calculate kinetic energy
546    kin_tot = np.zeros(nmolecs)
547    kin_tra = np.zeros(nmolecs)
548    kin_rni = np.zeros(nmolecs)
549    kin_rot = np.zeros(nmolecs)
550    kin_int = np.zeros(nmolecs)
551    # loop over molecules
552    for idx_molec, (idx_atm_init, idx_atm_end) in enumerate(zip(molec_idx[:-1], molec_idx[1:])):
553        # if monoatomic molecule
554        if idx_atm_end == idx_atm_init + 1:
555            v = vel[idx_atm_init]
556            m = masses[idx_atm_init]
557            kin_tot[idx_molec] = .5 * m * np.dot(v, v)
558            kin_tra[idx_molec] = kin_tot[idx_molec]
559            kin_rni[idx_molec] = 0
560            kin_rot[idx_molec] = 0
561            kin_int[idx_molec] = 0
562            continue
563        # compute center of mass position, velocity and total mass
564        com_r = np.zeros(3)
565        com_v = np.zeros(3)
566        com_m = 0
567        # loop over atoms in molecule
568        for r, v, m in zip(pos[idx_atm_init:idx_atm_end],
569                           vel[idx_atm_init:idx_atm_end],
570                           masses[idx_atm_init:idx_atm_end]):
571            com_r += m*r
572            com_v += m*v
573            com_m += m
574
575            # total kinetic energy is straightforward
576            kin_tot[idx_molec] += .5 * m * np.dot(v, v)
577
578        com_r /= com_m
579        com_v /= com_m
580
581        # translational kinetic energy
582        kin_tra[idx_molec] = .5 * com_m * np.dot(com_v, com_v)
583        # combined rotational and internal kinetic energy
584        kin_rni[idx_molec] = kin_tot[idx_molec] - kin_tra[idx_molec]
585
586        # compute tensor of inertia and angular momentum
587        inertia = np.zeros((3, 3))
588        angular_mom = np.zeros(3)
589        # loop over atoms in molecule
590        for r, v, m in zip(pos[idx_atm_init:idx_atm_end],
591                           vel[idx_atm_init:idx_atm_end],
592                           masses[idx_atm_init:idx_atm_end]):
593            # relative positions and velocities
594            rr = r - com_r
595            rv = v - com_v
596            rr2 = np.dot(rr, rr)
597            # inertia tensor:
598            #   (i,i) = m*(r*r - r(i)*r(i))
599            #   (i,j) = m*r(i)*r(j) (i != j)
600            atm_inertia = -m*np.tensordot(rr, rr, axes=0)
601            for i in range(3):
602                atm_inertia[i][i] += m*rr2
603            inertia += atm_inertia
604            # angular momentum: r x p
605            angular_mom += m * np.cross(rr, rv)
606
607        # angular velocity of the molecule: inertia^{-1} * angular_mom
608        angular_v = np.dot(np.linalg.inv(inertia), angular_mom)
609
610        # test_kin = 0
611        # for r, v, m in zip(pos[idx_atm_init:idx_atm_end],
612        #                    vel[idx_atm_init:idx_atm_end],
613        #                    masses[idx_atm_init:idx_atm_end]):
614        #     # relative positions and velocities
615        #     rr = r - com_r
616        #     rv = v - com_v - np.cross(angular_v, rr)
617        #     test_kin += .5 * m * np.dot(rv, rv)
618        #
619        # print(test_kin)
620
621        kin_rot[idx_molec] = .5 * np.dot(angular_v, angular_mom)
622        kin_int[idx_molec] = kin_rni[idx_molec] - kin_rot[idx_molec]
623
624        # end loop over molecules
625
626    return {'tot': kin_tot,
627            'tra': kin_tra,
628            'rni': kin_rni,
629            'rot': kin_rot,
630            'int': kin_int}
631
632
633def group_kinetic_energy(kin_molec, nmolecs, molec_group=None):
634    r"""
635    Sums up the partitioned kinetic energy for a
636    given group or the entire system.
637
638    Parameters
639    ----------
640    kin_molec : List[dict]
641        Partitioned kinetic energies per molecule.
642    nmolecs : int
643        Total number of molecules in the system.
644    molec_group : iterable
645        Indeces of the group to be summed up. None defaults to all molecules
646        in the system. Default: None.
647
648    Returns
649    -------
650    kin : dict
651        Dictionary of partitioned kinetic energy for the group.
652    """
653    #
654    kin = {'tot': 0, 'tra': 0, 'rni': 0, 'rot': 0, 'int': 0}
655    #
656    if molec_group is None:
657        molec_group = range(nmolecs)
658    # loop over molecules
659    for idx_molec in molec_group:
660        for key in kin.keys():
661            kin[key] += kin_molec[key][idx_molec]
662
663    return kin
664
665
666def group_ndof(ndof_molec, nmolecs, molec_group=None):
667    r"""
668    Sums up the partitioned degrees of freedom for a
669    given group or the entire system.
670
671    Parameters
672    ----------
673    ndof_molec : List[dict]
674        Partitioned degrees of freedom per molecule.
675    nmolecs : int
676        Total number of molecules in the system.
677    molec_group : iterable
678        Indeces of the group to be summed up. None defaults to all molecules
679        in the system. Default: None.
680
681    Returns
682    -------
683    ndof : dict
684        Dictionary of partitioned degrees of freedom for the group.
685    """
686    #
687    ndof = {'tot': 0, 'tra': 0, 'rni': 0, 'rot': 0, 'int': 0}
688    #
689    if molec_group is None:
690        molec_group = range(nmolecs)
691    # loop over molecules
692    for idx_molec in molec_group:
693        for key in ndof.keys():
694            ndof[key] += ndof_molec[idx_molec][key]
695
696    return ndof
697
698
699def calc_temperatures(kin_molec, ndof_molec, nmolecs, molec_group=None):
700    r"""
701    Calculates the partitioned temperature for a
702    given group or the entire system.
703
704    Parameters
705    ----------
706    kin_molec : List[dict]
707        Partitioned kinetic energies per molecule.
708    ndof_molec : List[dict]
709        Partitioned degrees of freedom per molecule.
710    nmolecs : int
711        Total number of molecules in the system.
712    molec_group : iterable
713        Indeces of the group to be summed up. None defaults to all molecules
714        in the system. Default: None.
715
716    Returns
717    -------
718    temp : dict
719        Dictionary of partitioned temperatures for the group.
720    """
721
722    kin = group_kinetic_energy(kin_molec, nmolecs, molec_group)
723    ndof = group_ndof(ndof_molec, nmolecs, molec_group)
724
725    temp = {}
726    for key in kin:
727        temp[key] = temperature(kin[key], ndof[key])
728
729    return temp
730
731
732def test_mb_dist(kin_molec, ndof_molec, nmolecs,
733                 temp, alpha, dict_keys, group=None,
734                 verbosity=0, screen=False, filename=None,
735                 ene_unit=None):
736    r"""
737    Tests if the partitioned kinetic energy trajectory of a group (or,
738    if group is None, of the entire system) are separately Maxwell-Boltzmann
739    distributed.
740
741    Parameters
742    ----------
743    kin_molec : List[List[dict]]
744        Partitioned kinetic energies per molecule for every frame.
745    ndof_molec : List[dict]
746        Partitioned degrees of freedom per molecule.
747    nmolecs : int
748        Total number of molecules in the system.
749    temp : float
750        Target temperature of the simulation.
751    alpha : float
752        Confidence for Maxwell-Boltzmann test.
753    dict_keys : List[str]
754        List of dictionary keys representing the partitions of the degrees
755        of freedom.
756    group : iterable
757        Indeces of the group to be tested. None defaults to all molecules
758        in the system. Default: None.
759    verbosity : int
760        Verbosity level, where 0 is quiet and 3 very chatty. Default: 0.
761    screen : bool
762        Plot distributions on screen. Default: False.
763    filename : string
764        Plot distributions to `filename`.pdf. Default: None.
765    ene_unit : string
766        Energy unit - used for plotting only.
767
768    Returns
769    -------
770    result : List[float]
771        p value for every partition
772    """
773    # save the partitioned kinetic energy trajectories
774    group_kin = {key: [] for key in dict_keys}
775    ndof = group_ndof(ndof_molec, nmolecs, group)
776    # loop over frames
777    for k in kin_molec:
778        frame_kin = group_kinetic_energy(k, nmolecs, group)
779        for key in dict_keys:
780            group_kin[key].append(frame_kin[key])
781
782    result = []
783    failed = 0
784    # test tot, tra, rni, rot, int
785    if verbosity > 1:
786        print('Testing whether\n'
787              '* total,\n'
788              '* translational,\n'
789              '* rotational & internal,\n'
790              '* rotational, and\n'
791              '* internal\n'
792              'kinetic energies are Maxwell-Boltzmann distributed.')
793    elif verbosity > 0:
794        print('Testing whether kinetic energies are Maxwell-Boltzmann distributed.')
795
796    for key in dict_keys:
797        p = check_mb_ensemble(kin=group_kin[key], temp=temp, ndof=ndof[key],
798                              alpha=alpha, verbosity=verbosity > 2,
799                              screen=screen, filename=filename+'_'+key,
800                              ene_unit=ene_unit)
801        result.append(p)
802        if alpha is not None and p < alpha:
803            failed += 1
804        if verbosity > 1 and alpha is not None:
805            if p >= alpha:
806                print('* {}: passed'.format(key))
807            else:
808                print('* {}: failed'.format(key))
809
810    if verbosity > 0 and alpha is not None:
811        if failed == 0:
812            print('-> Passed')
813        else:
814            print('-> Failed')
815
816    return result
817
818
819def test_temp_diff(kin_molec, ndof_molec, nmolecs,
820                   dtemp, dict_keys, group=None,
821                   verbosity=0, screen=False, filename=None,
822                   ene_unit=None):
823    r"""
824    Tests if the partitioned temperatures (averaged over a trajectory)
825    of a group (or, if group is None, of the entire system) are within a
826    range `dtemp` of the total temperature.
827
828    Parameters
829    ----------
830    kin_molec : List[List[dict]]
831        Partitioned kinetic energies per molecule for every frame.
832    ndof_molec : List[dict]
833        Partitioned degrees of freedom per molecule.
834    nmolecs : int
835        Total number of molecules in the system.
836    dtemp : float
837        Target temperature of the simulation.
838    dict_keys : List[str]
839        List of dictionary keys representing the partitions of the degrees
840        of freedom.
841    group : iterable
842        Indeces of the group to be tested. None defaults to all molecules
843        in the system. Default: None.
844    verbosity : int
845        Verbosity level, where 0 is quiet and 3 very chatty. Default: 0.
846    screen : bool
847        Plot distributions on screen. Default: False.
848    filename : string
849        Plot distributions to `filename`.pdf. Default: None.
850    ene_unit : string
851        Energy unit - used for plotting only.
852
853    Returns
854    -------
855    result : List[float]
856        Temperature ratio to the total temperature for every partition.
857    """
858    # save the partitioned temperature trajectories
859    group_temp = {key: [] for key in dict_keys}
860    # loop over frames
861    for k in kin_molec:
862        frame_temp = calc_temperatures(k, ndof_molec, nmolecs, group)
863        for key in dict_keys:
864            group_temp[key].append(frame_temp[key])
865    # average temperature
866    group_temp_avg = {}
867    for key in dict_keys:
868        group_temp_avg[key] = np.mean(group_temp[key])
869
870    failed = 0
871    result = []
872    if verbosity > 0 and dtemp is not None:
873        print('Testing whether temperatures')
874        print('  ' + str(dict_keys[1:]))
875        print('are within {:.1f}% of temperature'.format(dtemp*100))
876        print('  ' + dict_keys[0] + ' (' + str(group_temp_avg[dict_keys[0]]) + ')')
877    elif verbosity > 0:
878        print('Testing difference between temperatures')
879        print('  ' + str(dict_keys[1:]))
880        print('and')
881        print('  ' + dict_keys[0] + ' (' + str(group_temp_avg[dict_keys[0]]) + ')')
882
883    temp0 = group_temp_avg[dict_keys[0]]
884    for key in dict_keys[1:]:
885        result.append(group_temp_avg[key] / temp0)
886        if dtemp is not None:
887            if (1 - dtemp) * temp0 <= group_temp_avg[key] <= (1 + dtemp) * temp0:
888                if verbosity > 1:
889                    print('* {} ({:f}): passed'.format(key, group_temp_avg[key]))
890            else:
891                failed += 1
892                if verbosity > 1:
893                    print('* {} ({:f}): failed'.format(key, group_temp_avg[key]))
894
895    if verbosity > 0 and dtemp is not None:
896        if failed == 0:
897            print('-> Passed')
898        else:
899            print('-> Failed')
900
901    do_plot = screen or filename is not None
902    if do_plot:
903        data = []
904        for key in dict_keys:
905            t = group_temp[key]
906            data.append({'x': range(0, len(t)),
907                         'y': t,
908                         'name': 'T(' + key + ')'})
909
910        unit = ''
911        if ene_unit is not None:
912            unit = ' [' + ene_unit + ']'
913        plot.plot(data,
914                  legend='best',
915                  title='Temperature trajectories',
916                  xlabel='Frames',
917                  ylabel='Temperature' + unit,
918                  sci_x=True,
919                  filename=filename,
920                  screen=screen)
921
922    return result
923
924
925def test_temp_diff_groups(kin_molec, ndof_molec, nmolecs,
926                          group1, group2,
927                          dtemp, dict_keys, verbosity=0):
928    r"""
929    Tests if the partitioned temperatures (averaged over a trajectory)
930    of two groups are (individually) within a range `dtemp` of each others.
931
932    Parameters
933    ----------
934    kin_molec : List[List[dict]]
935        Partitioned kinetic energies per molecule for every frame.
936    ndof_molec : List[dict]
937        Partitioned degrees of freedom per molecule.
938    nmolecs : int
939        Total number of molecules in the system.
940    group1 : iterable
941        Indeces of the first group to be compared.
942    group2 : iterable
943        Indeces of the second group to be compared.
944    dtemp : float
945        Target temperature of the simulation.
946    dict_keys : List[str]
947        List of dictionary keys representing the partitions of the degrees
948        of freedom.
949    verbosity : int
950        Verbosity level, where 0 is quiet and 3 very chatty. Default: 0.
951
952    Returns
953    -------
954    result : List[float]
955        Temperature ratio (first group / second group) for every partition.
956    """
957    # save the partitioned temperature trajectories (group1)
958    group1_temp = {key: [] for key in dict_keys}
959    # loop over frames
960    for k in kin_molec:
961        frame_temp = calc_temperatures(k, ndof_molec, nmolecs, group1)
962        for key in dict_keys:
963            group1_temp[key].append(frame_temp[key])
964    # average temperature
965    for key in dict_keys:
966        group1_temp[key] = np.mean(group1_temp[key])
967
968    # save the partitioned temperature trajectories (group2)
969    group2_temp = {key: [] for key in dict_keys}
970    # loop over frames
971    for k in kin_molec:
972        frame_temp = calc_temperatures(k, ndof_molec, nmolecs, group2)
973        for key in dict_keys:
974            group2_temp[key].append(frame_temp[key])
975    # average temperature
976    for key in dict_keys:
977        group2_temp[key] = np.mean(group2_temp[key])
978
979    result = []
980    failed = 0
981    if verbosity > 0 and dtemp is not None:
982        print('Testing whether temperatures of both groups are within {:.1f}%'.
983              format(dtemp*100))
984
985    for key in dict_keys:
986        result.append(group1_temp[key]/group2_temp[key])
987        if dtemp is not None:
988            if (1. - dtemp) <= group1_temp[key]/group2_temp[key] <= (1 + dtemp):
989                if verbosity > 1:
990                    print('* {}: passed'.format(key))
991            else:
992                failed += 1
993                if verbosity > 1:
994                    print('* {}: failed'.format(key))
995
996    if verbosity > 0 and dtemp is not None:
997        if failed == 0:
998            print('-> Passed')
999        else:
1000            print('-> Failed')
1001
1002    return result
1003