request for help about autode HOT 2 CLOSED

import os.path
from typing import List
import matplotlib
import multiprocessing as mp

import numpy as np

import autode as ade
from autode.hessians import Hessian
from autode.values import Coordinates, Gradient, PotentialEnergy
from autode.wrappers.keywords import OptKeywords, SinglePointKeywords, KeywordsSet
from autode.opt.optimisers.base import BaseOptimiser
from autode.wrappers.methods import ExternalMethod
from autode.exceptions import NotImplementedInMethod
from autode.pes import UnRelaxedPES1D, RelaxedPESnD

import ase
from ase.optimize import BFGS
from ase.constraints import FixBondLengths

import torchani


class ANI(ExternalMethod):

    def __init__(self, ase_calculator, *args, **kwargs):
        super(ANI, self).__init__(*args, **kwargs)
        self.ase_calculator = ase_calculator

    def execute(self, calc: "CalculationExecutor") -> None:

        if isinstance(calc.input.keywords, OptKeywords):
            self._execute_optimisation(calc)
        elif isinstance(calc.input.keywords, SinglePointKeywords):
            pass  # calculated in _energy_from
        else:
            raise NotImplementedError

        open("out", "w").close()
        return None

    def _execute_optimisation(self, calc: "CalculationExecutor") -> None:
        """Execute an optimisation on the atoms"""

        ase_atoms = self._ase_atoms_from_molecule_in(calc)

        if calc.molecule.constraints.distance is not None:
            constraints = calc.molecule.constraints.distance
            print(list(constraints.values()))
            ase_atoms.set_constraint(FixBondLengths(
                pairs=list(constraints.keys()),
                bondlengths=list(constraints.values()),
                tolerance=1E-6
            ))

        opt = BFGS(ase_atoms)
        opt.run(fmax=1E-4)
        print("Optimisation done!")

        calc.molecule.coordinates = ase_atoms.positions

        return None

    def _ase_atoms_from_molecule_in(self, calc: "CalculationExecutor"):
        """Get a set of ASE atoms with a set calculator"""

        atoms = ase.Atoms(symbols=calc.molecule.atomic_symbols,
                          positions=calc.molecule.coordinates)
        atoms.set_calculator(self.ase_calculator)

        return atoms

    def implements(self,
                   calculation_type: "autode.calculations.types.CalculationType") -> bool:
        return True  # TODO: be more granular than this

    def hessian_from(self,
                     calc: "autode.calculations.executors.CalculationExecutor") -> Hessian:
        raise NotImplementedInMethod

    def terminated_normally_in(self, calc: "CalculationExecutor") -> bool:
        return True

    class Optimiser(BaseOptimiser):
        """A not very realistic optimiser"""

        @property
        def converged(self) -> bool:
            return True

        @property
        def last_energy_change(self) -> PotentialEnergy:
            return PotentialEnergy(0)

    def optimiser_from(self,
                       calc: "CalculationExecutor") -> "autode.opt.optimisers.base.BaseOptimiser":
        return self.Optimiser()

    def _energy_from(self, calc: "CalculationExecutor") -> PotentialEnergy:
        ase_atoms = self._ase_atoms_from_molecule_in(calc)
        return PotentialEnergy(ase_atoms.get_potential_energy(), units="eV")

    def gradient_from(self, calc: "CalculationExecutor") -> Gradient:
        ase_atoms = self._ase_atoms_from_molecule_in(calc)
        return Gradient(-ase_atoms.get_forces(), units="eV Å^-1")

    def coordinates_from(self, calc: "CalculationExecutor") -> Coordinates:
        return calc.molecule.coordinates

    def partial_charges_from(self, calc: "CalculationExecutor") -> List[float]:
        raise NotImplementedInMethod

    def version_in(self, calc: "CalculationExecutor") -> str:
        pass

    @staticmethod
    def input_filename_for(calc: "CalculationExecutor") -> str:
        return "in"

    @staticmethod
    def output_filename_for(calc: "CalculationExecutor") -> str:
        if os.path.exists("out"):
            os.remove("out")
        return "out"

    def generate_input_for(self, calc: "CalculationExecutor") -> None:
        return open("in", "w").close()

    def __repr__(self):
        return "ANI"


def plot_1d_unrelaxed():

    pes = UnRelaxedPES1D(h2o, rs={(0, 1): np.linspace(0.8, 1.5, num=20)})
    pes.calculate(method=ani)

    # matplotlib.use('TkAgg')   # <- I need to set this :shrug
    pes.plot("h2o_1d_unrelaxed.pdf")
    print(pes.relative_energies)


def plot_1d_relaxed():

    pes = RelaxedPESnD(h2o, rs={(0, 1): np.linspace(0.8, 1.5, num=20)})
    pes.calculate(method=ani)

    # ASE opt doesn't respect the fixed bond lengths?!
    print(pes.relative_energies)


if __name__ == '__main__':

    # Parallelism doesn't work..
    ade.Config.n_cores = 1
    mp.set_start_method("spawn", force=True)

    h2o = ade.Molecule(smiles="O")

    ani = ANI(
        ase_calculator=torchani.models.ANI1ccx().ase(),
        executable_name="python",
        keywords_set=KeywordsSet(),
        doi_list=[],
        implicit_solvation_type=None
    )

    # plot_1d_unrelaxed()
    plot_1d_relaxed()