#!/usr/bin/env python
# coding: utf-8

# ## Custom M3GNet model example (bulk Ru)
# Use a custom M3GNet model (trained by ParAMS or Simple Active Learning) by pointing to a parameter directory. The example below uses the RuH parameters for bulk Ru (hexagonal system).

from ase.build import bulk
from scm.base import ChemicalSystem
from scm.plams import view

ru = ChemicalSystem.from_ase_atoms(bulk("Ru", a=2.7, c=4.28))
print("Lattice angles: ", ru.lattice.get_angles(unit="degree"))
view(ru, direction="tilt_x", height=150, width=150, picture_path="picture1.png")


from scm.plams import Settings, AMSJob

s2 = Settings()
s2.runscript.nproc = 1
s2.input.ams.Task = "GeometryOptimization"
s2.input.ams.GeometryOptimization.OptimizeLattice = "Yes"
s2.input.ams.GeometryOptimization.Convergence.Quality = "Basic"
s2.input.MLPotential.Model = "Custom"
s2.input.MLPotential.Backend = "m3gnet"
# ams can handle environment variables when parsing the input - no need to expand it here
s2.input.MLPotential.ParameterDir = "$AMSHOME/atomicdata/MLPotential/M3GNet/RuH"

job2 = AMSJob(settings=s2, molecule=ru, name="m3gnet_custom_ru")
job2.run()

energy2 = job2.results.get_energy(unit="eV")
print(f"Custom M3GNet (RuH params): final energy {energy2:.3f} eV")


optimized_ru = job2.results.get_main_system()
density = optimized_ru.density(unit="g/cm3")
print(f"{density=:.2f} g/cm3")


# A lattice optimization may introduce a small amount of noise in the lattice:

print(optimized_ru.lattice.get_angles(unit="degree"))


# For symmetric systems it's a good idea to symmetrize the cell after a lattice optimization:

optimized_ru.symmetrize_cell()
print(optimized_ru.lattice.get_angles(unit="degree"))