Skip to content

Commit

Permalink
fix multilevel MBE
Browse files Browse the repository at this point in the history
  • Loading branch information
@alenaizan
alenaizan committed Feb 1, 2019
1 parent c76849e commit 4448c8a
Show file tree
Hide file tree
Showing 3 changed files with 121 additions and 173 deletions.
29 changes: 23 additions & 6 deletions psi4/driver/driver_nbody.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
from psi4 import core
from psi4.driver import p4util
from psi4.driver import constants
from psi4.driver import driver_nbody_multilevel
from psi4.driver.p4util.exceptions import *

from psi4.driver.task_base import BaseTask, unnp, plump_qcvar
Expand Down Expand Up @@ -586,14 +587,26 @@ def build_tasks(self, obj, bsse_type="all", **kwargs):
import json
template = json.dumps(kwargs)

# Get compute list
compute_dict = build_nbody_compute_list(self.bsse_type, self.max_nbody, self.max_frag)
# Store tasks by level
level = kwargs.pop('max_nbody', self.max_nbody)

# Add supersystem computation if requested
if level == 'supersystem':
data = json.loads(template)
data["molecule"] = self.molecule
self.task_list[str(level) + '_' + str(self.max_frag)] = obj(**data)
compute_dict = build_nbody_compute_list(self.bsse_type, self.max_nbody, self.max_frag)

else:
# Get compute list
compute_dict = build_nbody_compute_list(self.bsse_type, level, self.max_frag)

compute_list = compute_dict[bsse_type]

counter = 0
for count, n in enumerate(compute_list):
for key, pair in enumerate(compute_list[n]):
if pair in self.task_list:
if (str(level) + '_' + str(pair)) in self.task_list:
continue
print(pair)
data = json.loads(template)
Expand All @@ -607,7 +620,7 @@ def build_tasks(self, obj, bsse_type="all", **kwargs):
charges.extend([[chg, i] for i, chg in zip(positions, self.embedding_charges[frag])])
data['keywords'].update({'embedding_charges': charges})

self.task_list[pair] = obj(**data)
self.task_list[str(level) + '_' + str(pair)] = obj(**data)
counter += 1


Expand All @@ -631,10 +644,14 @@ def compute(self):
if self.embedding_charges:
core.set_global_option_python('EXTERN', None)

def prepare_results(self):
def prepare_results(self, results={}):

nlevels = len({i.split('_')[0] for i in self.task_list})
if nlevels > 1 and not results:
return driver_nbody_multilevel.prepare_results(self)

metadata = self.dict().copy()
results_list = {k: v.get_results() for k, v in self.task_list.items()}
results_list = {eval(k.split('_')[1]): v.get_results() for k, v in (results.items() or self.task_list.items())}
energies = {k: v['properties']["return_energy"] for k, v in results_list.items()}

ptype = None
Expand Down
226 changes: 82 additions & 144 deletions psi4/driver/driver_nbody_multilevel.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,161 +26,99 @@
# @END LICENSE
#

import pydantic

from typing import Dict, List, Any, Union

from psi4.driver import p4util
from psi4.driver.driver_nbody import NBodyComputer
from psi4.driver.task_base import BaseTask


class NBodyComputerMultiLevel(BaseTask):
driver: str
molecule: Any

# nbody multi-level kwargs
levels: Dict[Any, str] = {}
return_total_data: bool = False
return_wfn: bool = False

task_list: Dict[str, Any] = {}

def __init__(self, **data):
BaseTask.__init__(self, **data)

def build_tasks(self, packets, **kwargs):
supersystem = packets.pop('supersystem', None)

# Create NBodyComputers for each requested truncation order
for n, val in packets.items():
self.task_list[n] = NBodyComputer(**val['packet'], **kwargs, max_nbody=n)
val['packet'].pop('molecule', None)
self.task_list[n].build_tasks(val['computer'], **val['packet'])

# Create computers for calculations on the whole cluster if requested
if supersystem is not None:
self.task_list['supersystem'] = supersystem['computer'](**supersystem['packet'])
# Subtract MBE result of the highest requested MBE truncation order
n = max(packets)
self.task_list[str(n) + '_supersystem'] = NBodyComputer(**supersystem['packet'], **kwargs, max_nbody=n)
supersystem['packet'].pop('molecule', None)
self.task_list[str(n) + '_supersystem'].build_tasks(supersystem['computer'], **supersystem['packet'])


def plan(self):
ret = []
for k, v in self.task_list.items():
ret.append(v.plan())
return ret

def compute(self):
with p4util.hold_options_state():
# gof = core.get_output_file()
# core.close_outfile()
for k, v in self.task_list.items():
v.compute()

# core.set_output_file(gof, True)

def prepare_results(self):

from psi4.driver.driver_nbody import _print_nbody_energy

ptype = self.driver
natoms = self.molecule.natom()
bsse_type = list(self.task_list.values())[0].bsse_type
max_nbody = max([i for i in self.task_list if isinstance(i, int)])
supersystem = self.task_list.pop('supersystem', None)
n_supersystem = self.task_list.pop(str(max_nbody) + '_supersystem', None)

# Initialize with zeros
energy_result, gradient_result, hessian_result = 0, None, None
energy_body_contribution = {b: {} for b in bsse_type}
energy_body_dict = {b: {} for b in bsse_type}
if ptype in ['gradient', 'hessian']:
gradient_result = np.zeros((natoms, 3))
if ptype == 'hessian':
hessian_result = np.zeros((natoms * 3, natoms * 3))

for n in sorted(self.task_list)[::-1]:
result_n = self.task_list[n]
energy_bsse_dict = {b: 0 for b in result_n.bsse_type}
results = result_n.prepare_results()

for m in range(n - 1, n + 1):
if m == 0: continue
# Subtract the (n-1)-body contribution from the n-body contribution to get the n-body effect
sign = (-1)**(1 - m // n)
for b in result_n.bsse_type:
energy_bsse_dict[b] += sign * results['%s_energy_body_dict' %b.lower()]['%i%s' %(m, b.lower())]

if ptype == 'hessian':
hessian_result += sign * results['ptype_body_dict'][m]
gradient_result += sign * results['gradient_body_dict'][m]
if n == 1:
hessian1 = results['ptype_body_dict'][n]
gradient1 = results['gradient_body_dict'][n]

elif ptype == 'gradient':
gradient_result += sign * results['ptype_body_dict'][m]
# Keep 1-body contribution to compute interaction data
if n == 1:
gradient1 = results['ptype_body_dict'][n]

energy_result += energy_bsse_dict[bsse_type[0]]
for b in bsse_type:
energy_body_contribution[b][n] = energy_bsse_dict[b]

if supersystem:
# Super system recovers higher order effects at a lower level
supersystem_result = supersystem.get_results()
components = n_supersystem.prepare_results()

energy_result += supersystem_result['properties']['return_energy'] - components['energy_body_dict'][max_nbody]
for b in bsse_type:
energy_body_contribution[b][self.molecule.nfragments()] = (supersystem_result['properties']['return_energy'] -
components['energy_body_dict'][max_nbody])
def prepare_results(self):
"""
Use different levels of theory for different expansion levels
See kwargs description in driver_nbody.nbody_gufunc
:returns: *return type of func* |w--w| The data.
:returns: (*float*, :py:class:`~psi4.core.Wavefunction`) |w--w| data and wavefunction with energy/gradient/hessian set appropriately when **return_wfn** specified.
"""
from psi4.driver.driver_nbody import _print_nbody_energy

ptype = self.driver
natoms = self.molecule.natom()
supersystem = {k: v for k, v in self.task_list.items() if k.startswith('supersystem')}

# Initialize with zeros
energy_result, gradient_result, hessian_result = 0, None, None
energy_body_contribution = {b: {} for b in self.bsse_type}
energy_body_dict = {b: {} for b in self.bsse_type}
if ptype in ['gradient', 'hessian']:
gradient_result = np.zeros((natoms, 3))
if ptype == 'hessian':
hessian_result = np.zeros((natoms * 3, natoms * 3))

levels = {int(i.split('_')[0]) for i in self.task_list if 'supersystem' not in i}

for n in sorted(levels)[::-1]:
energy_bsse_dict = {b: 0 for b in self.bsse_type}
self.max_nbody = n
results = {k: v for k, v in self.task_list.items() if k.startswith(str(n))}
results = self.prepare_results(results=results)

for m in range(n - 1, n + 1):
if m == 0: continue
# Subtract the (n-1)-body contribution from the n-body contribution to get the n-body effect
sign = (-1)**(1 - m // n)
for b in self.bsse_type:
energy_bsse_dict[b] += sign * results['%s_energy_body_dict' %b.lower()]['%i%s' %(m, b.lower())]

if ptype == 'hessian':
gradient_result += supersystem_result['psi4:qcvars']['CURRENT GRADIENT'] - components['gradient_body_dict'][max_nbody]
hessian_result += supersystem_result['return_result'] - components['ptype_body_dict'][max_nbody]
hessian_result += sign * results['ptype_body_dict'][m]
gradient_result += sign * results['gradient_body_dict'][m]
if n == 1:
hessian1 = results['ptype_body_dict'][n]
gradient1 = results['gradient_body_dict'][n]

elif ptype == 'gradient':
gradient_result += supersystem_result['return_result'] - components['ptype_body_dict'][max_nbody]
gradient_result += sign * results['ptype_body_dict'][m]
# Keep 1-body contribution to compute interaction data
if n == 1:
gradient1 = results['ptype_body_dict'][n]

energy_result += energy_bsse_dict[self.bsse_type[0]]
for b in self.bsse_type:
energy_body_contribution[b][n] = energy_bsse_dict[b]

if supersystem:
# Super system recovers higher order effects at a lower level
supersystem_result = supersystem.pop('supersystem_' + str(self.max_frag)).get_results()
self.max_nbody = max(levels)
components = self.prepare_results(results=supersystem)

energy_result += supersystem_result['properties']['return_energy'] - components['energy_body_dict'][self.max_nbody]
for b in self.bsse_type:
energy_body_contribution[b][self.molecule.nfragments()] = (supersystem_result['properties']['return_energy'] -
components['energy_body_dict'][self.max_nbody])

if ptype == 'hessian':
gradient_result += supersystem_result['psi4:qcvars']['CURRENT GRADIENT'] - components['gradient_body_dict'][self.max_nbody]
hessian_result += supersystem_result['return_result'] - components['ptype_body_dict'][self.max_nbody]

for b in bsse_type:
for n in energy_body_contribution[b]:
energy_body_dict[b][n] = sum(
[energy_body_contribution[b][i] for i in range(1, n + 1) if i in energy_body_contribution[b]])

is_embedded = list(self.task_list.values())[0].embedding_charges
for b in bsse_type:
_print_nbody_energy(energy_body_dict[b], '%s-corrected multilevel many-body expansion' % b.upper(),
is_embedded)

if not self.return_total_data:
# Remove monomer cotribution for interaction data
energy_result -= energy_body_dict[bsse_type[0]][1]
if ptype in ['gradient', 'hessian']:
gradient_result -= gradient1
if ptype == 'hessian':
hessian_result -= hessian1
elif ptype == 'gradient':
gradient_result += supersystem_result['return_result'] - components['ptype_body_dict'][self.max_nbody]


energy_body_dict = {str(k) + b: v for b in energy_body_dict for k, v in energy_body_dict[b].items()}
for b in self.bsse_type:
for n in energy_body_contribution[b]:
energy_body_dict[b][n] = sum(
[energy_body_contribution[b][i] for i in range(1, n + 1) if i in energy_body_contribution[b]])

results = {'ret_energy': energy_result, 'ret_ptype': eval(ptype + '_result'),
'energy_body_dict': energy_body_dict, 'ptype_body_dict': {}}
is_embedded = self.embedding_charges
for b in self.bsse_type:
_print_nbody_energy(energy_body_dict[b], '%s-corrected multilevel many-body expansion' % b.upper(),
is_embedded)

return results
if not self.return_total_data:
# Remove monomer cotribution for interaction data
energy_result -= energy_body_dict[self.bsse_type[0]][1]
if ptype in ['gradient', 'hessian']:
gradient_result -= gradient1
if ptype == 'hessian':
hessian_result -= hessian1


def get_results(self):
# Use NBodyComputer functions as it has the same data structure
return NBodyComputer.get_results(self)
energy_body_dict = {str(k) + b: v for b in energy_body_dict for k, v in energy_body_dict[b].items()}

def get_psi_results(self, return_wfn=False):
# Use NBodyComputer functions as it has the same data structure
return NBodyComputer.get_psi_results(self, return_wfn)
return {'ret_energy': energy_result, 'ret_ptype': eval(ptype + '_result'), 'energy_body_dict': energy_body_dict, 'ptype_body_dict': {}}
39 changes: 16 additions & 23 deletions psi4/driver/task_planner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,14 +40,16 @@
from psi4.driver.task_base import SingleResult
from psi4.driver.driver_findif import FinDifComputer
from psi4.driver.driver_nbody import NBodyComputer
from psi4.driver.driver_nbody_multilevel import NBodyComputerMultiLevel
from psi4.driver.driver_cbs import CBSComputer, _cbs_text_parser
from psi4.driver.driver_util import negotiate_derivative_type

__all__ = ["task_planner"]


def _expand_cbs_methods(method, basis, driver, **kwargs):
if method == 'cbs' and kwargs.get('cbsmeta', None):
return method, basis, kwargs['cbsmeta']

# Expand CBS methods
if "/" in method:
kwargs["ptype"] = driver
Expand Down Expand Up @@ -118,34 +120,25 @@ def task_planner(driver, method, molecule, **kwargs):

# First check for BSSE type
if 'bsse_type' in kwargs:
# Multi-level MBE
if 'levels' in kwargs:
print('PLANNING Multi-level NBody')
plan = NBodyComputerMultiLevel(**packet, **kwargs)
packets = {}
for n, level in kwargs['levels'].items():
method, basis, cbsmeta = _expand_cbs_methods(level, basis, driver, **kwargs)
packets[n] = {'packet': packet.copy()}
packets[n]['packet'].update({'method': method, 'basis': basis, **cbsmeta})
packets[n]['computer'] = CBSComputer if method == 'cbs' else SingleResult

plan.build_tasks(packets, **kwargs)

return plan

# Single level MBE
else:
plan = NBodyComputer(**packet, **kwargs)
del packet["molecule"]
plan = NBodyComputer(**packet, **kwargs)
del packet["molecule"]

# Add tasks for every nbody level requested
levels = kwargs.pop('levels', {plan.max_nbody: method})
plan.max_nbody = max([i for i in levels if isinstance(i, int)])

for n, level in levels.items():
method, basis, cbsmeta = _expand_cbs_methods(level, basis, driver, cbsmeta=cbsmeta, **kwargs)
packet.update({'method': method, 'basis': basis})

if method == "cbs":
print('PLANNING NBody(CBS)')
plan.build_tasks(CBSComputer, **packet, **cbsmeta)
plan.build_tasks(CBSComputer, **packet, **cbsmeta, max_nbody=n)
else:
print('PLANNING NBody()')
plan.build_tasks(SingleResult, **packet)
plan.build_tasks(SingleResult, **packet, max_nbody=n)

return plan
return plan

# Check for CBS
elif method == "cbs":
Expand Down

0 comments on commit 4448c8a

Please sign in to comment.