11# -*- coding: utf-8 -*-
2- from __future__ import print_function
2+ from __future__ import print_function , annotations
33
44"""
55This module performs the relax over more
1515import sscha .Utilities as Utilities
1616import cellconstructor as CC
1717import cellconstructor .symmetries
18+ from sscha .aiida_ensemble import AiiDAEnsemble
1819
1920import sys , os
2021
4950
5051class SSCHA (object ):
5152
52- def __init__ (self , minimizer = None , ase_calculator = None , N_configs = 1 , max_pop = 20 ,
53- save_ensemble = False , cluster = None , ** kwargs ):
53+ def __init__ (
54+ self ,
55+ minimizer = None ,
56+ ase_calculator = None ,
57+ aiida_inputs : dict | None = None ,
58+ N_configs = 1 ,
59+ max_pop = 20 ,
60+ save_ensemble = False ,
61+ cluster = None ,
62+ ** kwargs
63+ ):
5464 """
5565 This module initialize the relaxer. It may perform
5666 constant volume or pressure relaxation using fully anharmonic potentials.
@@ -63,6 +73,16 @@ def __init__(self, minimizer = None, ase_calculator=None, N_configs=1, max_pop =
6373 ase_calculator : ase.calculators...
6474 An initialized ASE calculator. This will be used to compute energies and forces
6575 for the relaxation of the SSCHA.
76+ aiida_input: dict
77+ Dictionary containing the input data for the
78+ :class:`~sscha.aiida_ensemble.AiiDAEnsemble.compute_ensmble`
79+ method; namely:
80+ * pw_code: str,
81+ * protocol: str = 'moderate',
82+ * options: dict = None,
83+ * overrides: dict = None,
84+ * group_label: str = None,
85+ * kwargs
6686 N_configs : int
6787 The number of configuration to be used for each population
6888 max_pop: int, optional
@@ -84,6 +104,7 @@ def __init__(self, minimizer = None, ase_calculator=None, N_configs=1, max_pop =
84104 self .N_configs = N_configs
85105 self .max_pop = max_pop
86106 self .cluster = cluster
107+ self .aiida_inputs = aiida_inputs
87108 self .start_pop = 1
88109
89110 # If the ensemble must be saved at each iteration.
@@ -351,8 +372,11 @@ def relax(self, restart_from_ens = False, get_stress = False,
351372 self .minim .ensemble .generate (self .N_configs , sobol = sobol , sobol_scramble = sobol_scramble , sobol_scatter = sobol_scatter )
352373
353374 # Compute energies and forces
354- self .minim .ensemble .compute_ensemble (self .calc , get_stress ,
355- cluster = self .cluster )
375+ if isinstance (self .minim .ensemble , AiiDAEnsemble ):
376+ self .minim .ensemble .compute_ensemble (** self .aiida_inputs )
377+ else :
378+ self .minim .ensemble .compute_ensemble (
379+ self .calc , get_stress , cluster = self .cluster )
356380 #self.minim.ensemble.get_energy_forces(self.calc, get_stress)
357381
358382 if ensemble_loc is not None and self .save_ensemble :
@@ -573,8 +597,11 @@ def vc_relax(self, target_press = 0, static_bulk_modulus = 100,
573597 # self.minim.ensemble.save_bin(ensemble_loc, pop)
574598
575599 # Compute energies and forces
576- self .minim .ensemble .compute_ensemble (self .calc , True , stress_numerical ,
577- cluster = self .cluster )
600+ if isinstance (self .minim .ensemble , AiiDAEnsemble ):
601+ self .minim .ensemble .compute_ensemble (** self .aiida_inputs )
602+ else :
603+ self .minim .ensemble .compute_ensemble (
604+ self .calc , True , stress_numerical , cluster = self .cluster )
578605 #self.minim.ensemble.get_energy_forces(self.calc, True, stress_numerical = stress_numerical)
579606
580607 print ("RELAX force length:" , len (self .minim .ensemble .force_computed ))
0 commit comments