1# -*- coding: utf-8 -*-
2#
3# Copyright (c) 2020, the cclib development team
4#
5# This file is part of cclib (http://cclib.github.io) and is distributed under
6# the terms of the BSD 3-Clause License.
7
8"""Calculate properties of nuclei based on data parsed by cclib."""
9
10import logging
11
12import numpy as np
13
14from cclib.method.calculationmethod import Method
15from cclib.parser.utils import PeriodicTable
16from cclib.parser.utils import find_package
17
18_found_periodictable = find_package("periodictable")
19if _found_periodictable:
20    import periodictable as pt
21
22_found_scipy = find_package("scipy")
23if _found_scipy:
24    import scipy.constants
25
26
27def _check_periodictable(found_periodictable):
28    if not _found_periodictable:
29        raise ImportError("You must install `periodictable` to use this function")
30
31
32def _check_scipy(found_scipy):
33    if not _found_scipy:
34        raise ImportError("You must install `scipy` to use this function")
35
36
37def get_most_abundant_isotope(element):
38    """Given a `periodictable` element, return the most abundant
39    isotope.
40    """
41    most_abundant_isotope = element.isotopes[0]
42    abundance = 0
43    for iso in element:
44        if iso.abundance > abundance:
45            most_abundant_isotope = iso
46            abundance = iso.abundance
47    return most_abundant_isotope
48
49
50def get_isotopic_masses(charges):
51    """Return the masses for the given nuclei, respresented by their
52    nuclear charges.
53    """
54    _check_periodictable(_found_periodictable)
55    masses = []
56    for charge in charges:
57        el = pt.elements[charge]
58        isotope = get_most_abundant_isotope(el)
59        mass = isotope.mass
60        masses.append(mass)
61    return np.array(masses)
62
63
64class Nuclear(Method):
65    """A container for methods pertaining to atomic nuclei."""
66
67    def __init__(self, data, progress=None, loglevel=logging.INFO, logname="Log"):
68
69        self.required_attrs = ('natom','atomcoords','atomnos','charge')
70
71        super(Nuclear, self).__init__(data, progress, loglevel, logname)
72
73    def __str__(self):
74        """Return a string representation of the object."""
75        return "Nuclear"
76
77    def __repr__(self):
78        """Return a representation of the object."""
79        return "Nuclear"
80
81    def stoichiometry(self):
82        """Return the stoichemistry of the object according to the Hill system"""
83        cclib_pt = PeriodicTable()
84        elements = [cclib_pt.element[ano] for ano in self.data.atomnos]
85        counts = {el: elements.count(el) for el in set(elements)}
86
87        formula = ""
88        elcount = lambda el, c: "%s%i" % (el, c) if c > 1 else el
89        if 'C' in elements:
90            formula += elcount('C', counts['C'])
91            counts.pop('C')
92            if 'H' in elements:
93              formula += elcount('H', counts['H'])
94              counts.pop('H')
95        for el, c in sorted(counts.items()):
96            formula += elcount(el, c)
97
98        if getattr(self.data, 'charge', 0):
99            magnitude = abs(self.data.charge)
100            sign = "+" if self.data.charge > 0 else "-"
101            formula += "(%s%i)" % (sign, magnitude)
102        return formula
103
104    def repulsion_energy(self, atomcoords_index=-1):
105        """Return the nuclear repulsion energy."""
106        nre = 0.0
107        for i in range(self.data.natom):
108            ri = self.data.atomcoords[atomcoords_index][i]
109            zi = self.data.atomnos[i]
110            for j in range(i+1, self.data.natom):
111                rj = self.data.atomcoords[0][j]
112                zj = self.data.atomnos[j]
113                d = np.linalg.norm(ri-rj)
114                nre += zi*zj/d
115        return nre
116
117    def center_of_mass(self, atomcoords_index=-1):
118        """Return the center of mass."""
119        charges = self.data.atomnos
120        coords = self.data.atomcoords[atomcoords_index]
121        masses = get_isotopic_masses(charges)
122
123        mwc = coords * masses[:, np.newaxis]
124        numerator = np.sum(mwc, axis=0)
125        denominator = np.sum(masses)
126
127        return numerator / denominator
128
129    def moment_of_inertia_tensor(self, atomcoords_index=-1):
130        """Return the moment of inertia tensor."""
131        charges = self.data.atomnos
132        coords = self.data.atomcoords[atomcoords_index]
133        masses = get_isotopic_masses(charges)
134
135        moi_tensor = np.empty((3, 3))
136
137        moi_tensor[0][0] = np.sum(masses * (coords[:, 1]**2 + coords[:, 2]**2))
138        moi_tensor[1][1] = np.sum(masses * (coords[:, 0]**2 + coords[:, 2]**2))
139        moi_tensor[2][2] = np.sum(masses * (coords[:, 0]**2 + coords[:, 1]**2))
140
141        moi_tensor[0][1] = np.sum(masses * coords[:, 0] * coords[:, 1])
142        moi_tensor[0][2] = np.sum(masses * coords[:, 0] * coords[:, 2])
143        moi_tensor[1][2] = np.sum(masses * coords[:, 1] * coords[:, 2])
144
145        moi_tensor[1][0] = moi_tensor[0][1]
146        moi_tensor[2][0] = moi_tensor[0][2]
147        moi_tensor[2][1] = moi_tensor[1][2]
148
149        return moi_tensor
150
151    def principal_moments_of_inertia(self, units='amu_bohr_2'):
152        """Return the principal moments of inertia in 3 kinds of units:
153        1. [amu][bohr]^2
154        2. [amu][angstrom]^2
155        3. [g][cm]^2
156        and the principal axes.
157        """
158        choices = ('amu_bohr_2', 'amu_angstrom_2', 'g_cm_2')
159        units = units.lower()
160        if units not in choices:
161            raise ValueError("Invalid units, pick one of {}".format(choices))
162        moi_tensor = self.moment_of_inertia_tensor()
163        principal_moments, principal_axes = np.linalg.eigh(moi_tensor)
164        if units == "amu_bohr_2":
165            _check_scipy(_found_scipy)
166            bohr2ang = scipy.constants.value("atomic unit of length") / scipy.constants.angstrom
167            conv = 1 / bohr2ang ** 2
168        elif units == "amu_angstrom_2":
169            conv = 1
170        elif units == "g_cm_2":
171            _check_scipy(_found_scipy)
172            amu2g = scipy.constants.value("unified atomic mass unit") * scipy.constants.kilo
173            conv = amu2g * (scipy.constants.angstrom / scipy.constants.centi) ** 2
174        return conv * principal_moments, principal_axes
175
176    def rotational_constants(self, units='ghz'):
177        """Compute the rotational constants in 1/cm or GHz."""
178        choices = ('invcm', 'ghz')
179        units = units.lower()
180        if units not in choices:
181            raise ValueError("Invalid units, pick one of {}".format(choices))
182        principal_moments = self.principal_moments_of_inertia("amu_angstrom_2")[0]
183        _check_scipy(_found_scipy)
184        bohr2ang = scipy.constants.value('atomic unit of length') / scipy.constants.angstrom
185        xfamu = 1 / scipy.constants.value('electron mass in u')
186        xthz = scipy.constants.value('hartree-hertz relationship')
187        rotghz = xthz * (bohr2ang ** 2) / (2 * xfamu * scipy.constants.giga)
188        if units == 'ghz':
189            conv = rotghz
190        if units == 'invcm':
191            ghz2invcm = scipy.constants.giga * scipy.constants.centi / scipy.constants.c
192            conv = rotghz * ghz2invcm
193        return conv / principal_moments
194
195
196del find_package
197