Incompatibility with CP2K Calculator

[potus28]

Hello,

I am trying to use Flare++ for on-the-fly training using CP2K as the DFT calculator. The code for my simulation is as follows:

# Import numpy and matplotlib
import numpy as np
from numpy.random import random
import matplotlib.pyplot as plt
import matplotlib

# flare++ imports
import flare_pp._C_flare as flare_pp
from flare_pp.sparse_gp import SGP_Wrapper
from flare_pp.sparse_gp_calculator import SGP_Calculator

# flare imports
from flare.ase.otf import ASE_OTF
from flare import otf_parser

# ASE imports
from ase import Atoms, units
from ase.build import supercells
from ase.visualize import view
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary, ZeroRotation
from ase.build import fcc111, add_adsorbate
from ase.io import write
from ase.io.trajectory import Trajectory
from ase.constraints import FixAtoms
from ase.calculators.vasp import Vasp

# Custom calculators
from mycalculators import *

# ASE calculator
calc = CP2KCalculator(ecut=400, functional="rVV10", scf=30, orbital_transform=True)

# Create the system
traj = Trajectory("MTD.traj")
atoms = traj[-1]
n_atoms = len(atoms)

# Freeze the bottom half of the slab
#c = FixAtoms(indices=[atom.index for atom in atoms if atom.position[2] < 0.5*atoms.get_cell()[2][2])
#atoms.set_constraint(c)

# Set the initial velocity
temperature = 450  # in K
MaxwellBoltzmannDistribution(atoms, temperature_K=temperature)
Stationary(atoms)  # zero linear momentum
ZeroRotation(atoms)  # zero angular momentum

# Set MD parameters.
md_engine = "Langevin"
md_dict = {"friction": 0.01, "temperature_K": temperature}

# Create sparse GP model.
species_map = {42: 0 , 6: 1, 1:2, 8:3}  # Molybdenum (atomic number 42) is species 0
cutoff = 5.0  # in A
sigma = 2.0  # in eV
power = 2  # power of the dot product kernel
kernel = flare_pp.NormalizedDotProduct(sigma, power)
cutoff_function = "quadratic"
many_body_cutoffs = [cutoff]
radial_basis = "chebyshev"
radial_hyps = [0., cutoff]
cutoff_hyps = []
n_species = 4
N = 12
lmax = 3
descriptor_settings = [n_species, N, lmax]
descriptor_calculator = flare_pp.B2(
  radial_basis,
  cutoff_function,
  radial_hyps,
  cutoff_hyps,
  descriptor_settings
)

# Set the noise values.
sigma_e = 0.001 * n_atoms  # eV (1 meV/atom)
sigma_f = 0.05  # eV/A
sigma_s = 0.0006  # eV/A^3 (about 0.1 GPa)

# Choose uncertainty type.
# Other options are "DTC" (Deterministic Training Conditional) or
# "SOR" (Subset of Regressors).
variance_type = "local"  # Compute uncertainties on local energies (normalized)

# Choose settings for hyperparameter optimization.
max_iterations = 20  # Max number of BFGS iterations during optimization
opt_method = "L-BFGS-B"  # Method used for hyperparameter optimization

# Bounds for hyperparameter optimization.
# Keeps the energy noise from going to zero.
bounds = [(None, None), (sigma_e, None), (None, None), (None, None)]

# Create a model wrapper that is compatible with the flare code.
gp_model = SGP_Wrapper([kernel], [descriptor_calculator], cutoff,
                    sigma_e, sigma_f, sigma_s, species_map,
                    variance_type=variance_type,
                    stress_training=False,
                    opt_method=opt_method,
                    bounds=bounds,
                    max_iterations=max_iterations)

# Create an ASE calculator based on the GP model.
flare_calculator = SGP_Calculator(gp_model)

# Set up OTF object.
init_atoms = list(range(n_atoms))  # Initial environments to include in the sparse set
output_name = 'flare'  # Name of the output file
std_tolerance_factor = -0.01  # Uncertainty tolerance for calling DFT
freeze_hyps = 2  # Freeze hyperparameter optimization after this many DFT calls
min_steps_with_model = 10  # Min number of steps between DFT calls
update_style = "threshold"  # Strategy for adding sparse environments
update_threshold = 0.005  # Threshold for determining which sparse environments to add

otf_params = {'init_atoms': init_atoms, 'output_name': output_name,
              'std_tolerance_factor': std_tolerance_factor,
              'freeze_hyps': freeze_hyps,
              'min_steps_with_model': min_steps_with_model,
              'update_style': update_style,
              'update_threshold': update_threshold}

# Create OTF object.
timestep = 0.5 * units.fs
number_of_steps = 2000
otf = ASE_OTF(
    atoms, timestep=timestep, number_of_steps=number_of_steps,
    dft_calc=calc, md_engine=md_engine, md_kwargs=md_dict,
    calculator=flare_calculator, **otf_params)

# Run on-the-fly dynamics.
otf.run()

And here are the contents of mycalculators.py. All this file does is help automate additional input with a CP2K calculator:

from ase.calculators.cp2k import CP2K
from ase.units import Rydberg

def CP2KCalculator(ecut, functional="LDA", kpoints=None, dipole_correction=False, 
                   orbital_transform=False,smearing=False, 
                   voronoi=False, cube=False, bqb=False,v_hartree=False, added_mos=None,scf=50, efield=None, stress=True, potential = "GTH-PBE"):
    """Creates a CP2K calculator object with Settings for Production Runs"""

    # By Default, assume we want to have the walltime as just shy of 48 hours
    inp = '''
&GLOBAL
WALLTIME 47:59:00
&END GLOBAL
&FORCE_EVAL
&DFT
'''

    ### DFT SECTION
    if dipole_correction:
        inp += '''
SURFACE_DIPOLE_CORRECTION .TRUE.
SURF_DIP_DIR Z  
'''
    if efield is not None:
        stress = False
        s = "INTENSITY " + str(efield) + "\n"
        inp += '''
&PERIODIC_EFIELD
'''
        inp += s
        inp += '''
&END PERIODIC_EFIELD
'''      
    if kpoints is not None:

        s = "SCHEME MONKHORST-PACK " + str(kpoints[0]) + " " + str(kpoints[1]) + " " + str(kpoints[2]) + "\n"
        inp += '''
&KPOINTS
'''
        inp += s
        inp += '''
&END KPOINTS
'''      

    ### SCF SECTION
    inp += '''
&SCF
&OUTER_SCF .TRUE.
MAX_SCF 50
&END OUTER_SCF
'''

    if orbital_transform:
        inp +='''
&OT .TRUE.
MINIMIZER DIIS
ALGORITHM IRAC
PRECONDITIONER FULL_SINGLE_INVERSE
&END OT
'''

    if smearing:
        if added_mos is not None:
            mos = added_mos
        else:
            mos = 10
        inp += "ADDED_MOS " + str(mos) + "\n"
        inp +='''
&SMEAR ON
METHOD FERMI_DIRAC
ELECTRONIC_TEMPERATURE [K] 300
&END SMEAR
&MIXING .TRUE.
METHOD BROYDEN_MIXING
&END MIXING
'''

    ###CLOSE SCF SECTION
    inp += '''
&END SCF
    '''

    ### XC Section
    inp += '''
&XC
&XC_GRID
XC_DERIV NN10_SMOOTH
XC_SMOOTH_RHO NN10
&END
'''
    if functional == "BLYP+D3":
        functional=None
        inp += '''
''''''
&XC_FUNCTIONAL BLYP
&END XC_FUNCTIONAL
&VDW_POTENTIAL
POTENTIAL_TYPE PAIR_POTENTIAL
&PAIR_POTENTIAL
R_CUTOFF 15.0
TYPE DFTD3
VERBOSE_OUTPUT
CALCULATE_C9_TERM .FALSE.
REFERENCE_FUNCTIONAL BLYP
PARAMETER_FILE_NAME dftd3.dat
&END PAIR_POTENTIAL
&END VDW_POTENTIAL
'''

    if functional == "PBE+D3":
        functional="PBE"
        inp += '''
''''''
&VDW_POTENTIAL
POTENTIAL_TYPE PAIR_POTENTIAL
&PAIR_POTENTIAL
R_CUTOFF 15.0
TYPE DFTD3
VERBOSE_OUTPUT
CALCULATE_C9_TERM .FALSE.
REFERENCE_FUNCTIONAL PBE
PARAMETER_FILE_NAME dftd3.dat
&END PAIR_POTENTIAL
&END VDW_POTENTIAL
'''

    if functional == "optB88-vdw":
        functional = None
        inp += '''
&XC_FUNCTIONAL

&GGA_X_OPTB88_VDW
&END GGA_X_OPTB88_VDW
&PW92
&END PW92
&END XC_FUNCTIONAL
&vdW_POTENTIAL
DISPERSION_FUNCTIONAL NON_LOCAL
&NON_LOCAL
TYPE DRSLL
VERBOSE_OUTPUT
KERNEL_FILE_NAME vdW_kernel_table.dat
&END NON_LOCAL
&END vdW_POTENTIAL
'''

    if functional == "optB86B-vdw":
        functional = None
        inp += '''
&XC_FUNCTIONAL
&GGA_X_OPTB86B_VDW
&END
&PW92
&END PW92
&END XC_FUNCTIONAL
&vdW_POTENTIAL
DISPERSION_FUNCTIONAL NON_LOCAL
&NON_LOCAL
TYPE DRSLL
VERBOSE_OUTPUT
KERNEL_FILE_NAME vdW_kernel_table.dat
&END NON_LOCAL
&END vdW_POTENTIAL
'''
    if functional == "optPBE-vdw":
        functional = None
        inp += '''
&XC_FUNCTIONAL
&GGA_X_OPTPBE_VDW
&END
&PW92
&END PW92
&END XC_FUNCTIONAL
&vdW_POTENTIAL
DISPERSION_FUNCTIONAL NON_LOCAL
&NON_LOCAL
TYPE DRSLL
VERBOSE_OUTPUT
KERNEL_FILE_NAME vdW_kernel_table.dat
&END NON_LOCAL
&END vdW_POTENTIAL
'''

    if functional == "rVV10":
        functional = None
        inp += '''
&XC_FUNCTIONAL
&GGA_X_RPW86
&END GGA_X_RPW86
&GGA_C_PBE
&END GGA_C_PBE
&END XC_FUNCTIONAL
&vdW_POTENTIAL
DISPERSION_FUNCTIONAL NON_LOCAL
&NON_LOCAL
TYPE RVV10
VERBOSE_OUTPUT
KERNEL_FILE_NAME rVV10_kernel_table.dat
&END NON_LOCAL
&END vdW_POTENTIAL
'''

    ### CLOSE OFF XC
    inp += '''
&END XC
'''
        ### DFT Print SECTION 
    if voronoi or cube or bqb or v_hartree:
        inp += '''
&PRINT
'''
        if cube:
            inp += '''
&E_DENSITY_CUBE
&END E_DENSITY_CUBE
'''
        if bqb:
            inp += '''
&E_DENSITY_BQB
&END E_DENSITY_BQB
'''
        if v_hartree:
            inp += '''
&V_HARTREE_CUBE
&END V_HARTREE_CUBE
'''
        if voronoi:
            inp += '''
&VORONOI ON
OUTPUT_TEXT .TRUE.
&END VORONOI
'''
        inp += '''
&END PRINT
'''

    #### CLOSE EVERYTHING 
    inp +='''
&END DFT
&END FORCE_EVAL
''' 

    calc = CP2K(
        auto_write=False,
        basis_set="DZVP-MOLOPT-SR-GTH",
        basis_set_file="BASIS_MOLOPT",
        charge=0,
        cutoff = ecut*Rydberg,
        debug = False,
        force_eval_method = "Quickstep",
        xc = functional,
        inp = inp,
        max_scf = scf,
        poisson_solver ="auto",
        potential_file = "POTENTIAL",
        pseudo_potential = potential,
        stress_tensor = stress,
        print_level = "LOW",
        uks=True
    )

    return calc

However, I run into this error at the beginning of the simulation on the first DFT call:

Traceback (most recent call last):
  File "/work/wnw36/TEST/FLARE++/adsorbate/run.py", line 132, in <module>
    otf.run()
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/site-packages/flare/otf.py", line 256, in run
    self.initialize_train()
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/site-packages/flare/ase/otf.py", line 174, in initialize_train
    super().initialize_train()
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/site-packages/flare/otf.py", line 357, in initialize_train
    self.run_dft()
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/site-packages/flare/otf.py", line 402, in run_dft
    forces = self.dft_module.run_dft_par(
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/site-packages/flare/ase/dft.py", line 29, in run_dft_par
    calc = deepcopy(dft_calc)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/work/wnw36/Programs/miniconda3/lib/python3.9/copy.py", line 161, in deepcopy
    rv = reductor(4)
TypeError: cannot pickle '_thread.lock' object
Exception ignored in: <function CP2K.__del__ at 0x2ad61f11c4c0>
Traceback (most recent call last):
  File "/home/wnw36/Programs/ase/dev/ase/calculators/cp2k.py", line 207, in __del__
    if self._shell:
AttributeError: 'CP2K' object has no attribute '_shell'
Exception ignored in: <function Cp2kShell.__del__ at 0x2ad61f11ca60>
Traceback (most recent call last):
  File "/home/wnw36/Programs/ase/dev/ase/calculators/cp2k.py", line 503, in __del__
    if self.isready:
AttributeError: 'Cp2kShell' object has no attribute 'isready'

I'm not sure what the error is, but it looks like something is going wrong with using CP2K as the DFT calculator. Is there anything I can change in the simulation so that I can use CP2K? Any help would be much appreciated. Thanks!

[YuuuXie]

Hi @potus28, thank you for your interest in our code. I'm not used the ASE CP2K calculator before. Here are a few suggestions on debugging this:

1. If you just set up a simple python script that uses ASE CP2K calculator to do a DFT calculation, with no flare-related things, will the same error occur? If so, this is then an issue with ASE CP2K calculator, for which you need to open an issue to ASE instead of FLARE.
2. If (1) works, maybe you can try the same settings of ASE CP2K calculator for FLARE, and if it breaks down you can provide us with a minimal reproducible script for debugging

[potus28]

Hi @YuuuXie , thank you so much for the reply and your suggestions. I double-checked to see if it was an issue with the ASE CP2K calculator, but it is working as expected. I also tried using just FLARE and the code failed with the same " TypeError: cannot pickle '_thread.lock' object" error. I've attached a python program based on the tutorial from the FLARE Read the Docs page with the built-in ASE CP2K calculator. Please let me know if you need anything else. Thank you so much!

# # Build the System with ASE

# In[43]:

import numpy as np
from ase import units
from ase.spacegroup import crystal
from ase.build import bulk
from ase.visualize import view

np.random.seed(12345)

a = 3.52678
super_cell = bulk("C", "diamond", a=a, cubic=True).repeat((2,2,2))

view(super_cell,viewer="x3d")

# # Set up the FLARE Calculator for OTF Gaussian Process Regression

# In[44]:

from flare.gp import GaussianProcess
from flare.utils.parameter_helper import ParameterHelper

# set up hyperparameters for a 2+3 body kernel
kernels = ["twobody", "threebody"]
parameters = {'cutoff_twobody': 5.0,
             "cutoff_threebody": 3.5}

pm = ParameterHelper(
    kernels = kernels,
    random = True,
    parameters = parameters
)

hm = pm.as_dict()
hyps = hm["hyps"]
cut = hm["cutoffs"]
print("hyps", hyps)

# In[38]:

# Since we are using ASE, always set compontent to "mc" regardless if the system is single component or not 
gp_model = GaussianProcess(
    kernels = kernels,
    component = "mc",
    hyps = hyps,
    cutoffs = cut,
    hyp_labels = ["sig2", "ls2", "sig3", "ls3", "noise"],
    opt_algorithm = "L-BFGS-B",
    ncpus = 1

)

# # Optional Mapped Gaussian Process

# If you want to use Mapped Gaussian Process (MGP), the set up is as follows

# In[39]:

from flare.mgp import MappedGaussianProcess

grid_params = {"twobody": {"grid_num":[64]},
               "threebody": {"grid_num":[16,16,16]}
              }

mgp_model = MappedGaussianProcess(
    grid_params,
    unique_species=[6],
    n_cpus=1,
    var_map=None
)

# # FLARE ASE Calculator

# In[40]:

from flare.ase.calculator import FLARE_Calculator
# If you want to use MGP, switch use_mapping to True and mgp_model to mgp_model
flare_calculator = FLARE_Calculator(
    gp_model,
    par = True,
    mgp_model = mgp_model ,
    use_mapping= True
)

# # DFT Calculator

# In[33]:

from ase.calculators.cp2k import CP2K
# ASE calculator
input_data = '''
&FORCE_EVAL
   &DFT
      &SCF
         &OUTER_SCF .TRUE.
            MAX_SCF 50
         &END OUTER_SCF
         &OT .TRUE.
            MINIMIZER DIIS
            ALGORITHM IRAC
            PRECONDITIONER FULL_SINGLE_INVERSE
         &END OT
      &END SCF
      &XC
         &XC_GRID
            XC_DERIV NN10_SMOOTH
            XC_SMOOTH_RHO NN10
         &END XC_GRID
      &END XC
   &END DFT
&END FORCE_EVAL
'''

dft_calc = CP2K(
        auto_write=False,
        basis_set="DZVP-MOLOPT-SR-GTH",
        basis_set_file="BASIS_MOLOPT",
        charge=0,
        cutoff = 400*units.Rydberg,
        debug = False,
        force_eval_method = "Quickstep",
        xc = "PBE",
        inp = input_data,
        max_scf = 30,
        poisson_solver ="auto",
        potential_file = "POTENTIAL",
        pseudo_potential = "GTH-PBE",
        stress_tensor = True,
        print_level = "LOW",
        uks=True
    )

# Double check to ensure the calculator is working

# In[47]:

super_cell.calc = dft_calc
en = super_cell.get_potential_energy()
print("Potential_energy (eV):", en)

# # OTF Molecular Dynamics Engine

# In[48]:

from ase import units
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary, ZeroRotation

temperature = 500
MaxwellBoltzmannDistribution(super_cell, temperature_K=temperature)

# Ensure that linear and angular momentum are zero
Stationary(super_cell)
ZeroRotation(super_cell)

md_engine = "VelocityVerlet"
md_kwargs = {}

#md_engine = "NVTBerendsen"
#md_kwargs = {"temperature": 500, "taut": 0.5e3 * units.fs}

# # Run MD with OFT Training!

# In[49]:

from flare.ase.otf import ASE_OTF

super_cell.calc = flare_calculator

otf_params = {'init_atoms': [0, 1, 2, 3],
              'output_name': 'otf',
              'std_tolerance_factor': 2,
              'max_atoms_added' : 4,
              'freeze_hyps': 10,
              'write_model': 3} # If you will probably resume the training, please set to 3

test_otf = ASE_OTF(super_cell,
                   timestep = 1 * units.fs,
                   number_of_steps = 3,
                   dft_calc = dft_calc,
                   md_engine = md_engine,
                   md_kwargs = md_kwargs,
                   **otf_params)

test_otf.run()

[YuuuXie]

Hi @potus28 , it seems that there is an issue with "deepcopy" the ASE CP2K calculator. I'm wondering, without flare, in your python script for using CP2K from ASE, could you try calc = deepcopy(cp2k_calc) before or after the calculation, and see if there is the same error? If so, there might be an issue with the CP2K calculator, and worth reporting to the ASE team.

[jiongwalai]

I have tried to modify the code of CP2K calculator in ASE (ase/calculators/cp2k.py), made it create shell before calculation and cancel shell after calculation. So that the function deepcopy(cp2k_calc) would not be conflicted with the running shell. The modified code is as follows:

200c200
<         #self._shell = Cp2kShell(self.command, self._debug)
---
>         self._shell = Cp2kShell(self.command, self._debug)
207,209c207,209
<         #if self._shell:
<         #    self._release_force_env()
<         #    del(self._shell)
---
>         if self._shell:
>             self._release_force_env()
>             del(self._shell)
245c245
<         self._shell = Cp2kShell(self.command, self._debug)
---
> 
313,315d312
< 
<         self._release_force_env()
<         del(self._shell)

[YuuuXie]

Hi @potus28, let me know if the above approach solves your problem!

And @jiongwalai, thank you very much for posting this patch! Maybe you can open an issue/pull request to ASE, letting them know about this problem, such that they can fix it in the new release.

[potus28]

HI @YuuuXie and @jiongwalai , sorry for the late response and I hope y'all have had a great holiday season. Thank y'all so much for the suggestions I will try this approach this week!

[potus28]

Hi @YuuuXie and @jiongwalai, following this approach, I was able to run a OTF simulation with ASE using CP2K as the DFT calculator. Thanks again for y'all's help!