1# -*- coding: utf-8 -*-
2#
3# Copyright (c) 2018, 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"""Generic file writer and related tools"""
9
10import logging
11from abc import ABC, abstractmethod
12from collections.abc import Iterable
13
14import numpy
15
16from cclib.parser.utils import PeriodicTable
17from cclib.parser.utils import find_package
18
19_has_openbabel = find_package("openbabel")
20if _has_openbabel:
21    from cclib.bridge import makeopenbabel
22    try:
23        from openbabel import openbabel as ob
24        import openbabel.pybel as pb
25    except:
26        import openbabel as ob
27        import pybel as pb
28
29
30class MissingAttributeError(Exception):
31    pass
32
33
34class Writer(ABC):
35    """Abstract class for writer objects."""
36
37    required_attrs = ()
38
39    def __init__(self, ccdata, jobfilename=None, indices=None, terse=False,
40                 *args, **kwargs):
41        """Initialize the Writer object.
42
43        This should be called by a subclass in its own __init__ method.
44
45        Inputs:
46          ccdata - An instance of ccData, parsed from a logfile.
47          jobfilename - The filename of the parsed logfile.
48          indices - One or more indices for extracting specific geometries/etc. (zero-based)
49          terse - Whether to print the terse version of the output file - currently limited to cjson/json formats
50        """
51
52        self.ccdata = ccdata
53        self.jobfilename = jobfilename
54        self.indices = indices
55        self.terse = terse
56        self.ghost = kwargs.get("ghost")
57
58        self.pt = PeriodicTable()
59
60        self._check_required_attributes()
61
62        # Open Babel isn't necessarily present.
63        if _has_openbabel:
64            # Generate the Open Babel/Pybel representation of the molecule.
65            # Used for calculating SMILES/InChI, formula, MW, etc.
66            self.obmol, self.pbmol = self._make_openbabel_from_ccdata()
67            self.bond_connectivities = self._make_bond_connectivity_from_openbabel(self.obmol)
68
69        self._fix_indices()
70
71    @abstractmethod
72    def generate_repr(self):
73        """Generate the written representation of the logfile data."""
74
75    def _calculate_total_dipole_moment(self):
76        """Calculate the total dipole moment."""
77
78        # ccdata.moments may exist, but only contain center-of-mass coordinates
79        if len(getattr(self.ccdata, 'moments', [])) > 1:
80            return numpy.linalg.norm(self.ccdata.moments[1])
81
82    def _check_required_attributes(self):
83        """Check if required attributes are present in ccdata."""
84        missing = [x for x in self.required_attrs
85                   if not hasattr(self.ccdata, x)]
86        if missing:
87            missing = ' '.join(missing)
88            raise MissingAttributeError(
89                'Could not parse required attributes to write file: ' + missing)
90
91    def _make_openbabel_from_ccdata(self):
92        """Create Open Babel and Pybel molecules from ccData."""
93        if not hasattr(self.ccdata, 'charge'):
94            logging.warning("ccdata object does not have charge, setting to 0")
95            _charge = 0
96        else:
97            _charge = self.ccdata.charge
98        if not hasattr(self.ccdata, 'mult'):
99            logging.warning("ccdata object does not have spin multiplicity, setting to 1")
100            _mult = 1
101        else:
102            _mult = self.ccdata.mult
103        obmol = makeopenbabel(self.ccdata.atomcoords,
104                              self.ccdata.atomnos,
105                              charge=_charge,
106                              mult=_mult)
107        if self.jobfilename is not None:
108            obmol.SetTitle(self.jobfilename)
109        return (obmol, pb.Molecule(obmol))
110
111    def _make_bond_connectivity_from_openbabel(self, obmol):
112        """Based upon the Open Babel/Pybel molecule, create a list of tuples
113        to represent bonding information, where the three integers are
114        the index of the starting atom, the index of the ending atom,
115        and the bond order.
116        """
117        bond_connectivities = []
118        for obbond in ob.OBMolBondIter(obmol):
119            bond_connectivities.append((obbond.GetBeginAtom().GetIndex(),
120                                        obbond.GetEndAtom().GetIndex(),
121                                        obbond.GetBondOrder()))
122        return bond_connectivities
123
124    def _fix_indices(self):
125        """Clean up the index container type and remove zero-based indices to
126        prevent duplicate structures and incorrect ordering when
127        indices are later sorted.
128        """
129        if not self.indices:
130            self.indices = set()
131        elif not isinstance(self.indices, Iterable):
132            self.indices = set([self.indices])
133        # This is the most likely place to get the number of
134        # geometries from.
135        if hasattr(self.ccdata, 'atomcoords'):
136            lencoords = len(self.ccdata.atomcoords)
137            indices = set()
138            for i in self.indices:
139                if i < 0:
140                    i += lencoords
141                indices.add(i)
142            self.indices = indices
143        return
144
145
146del find_package
147