def multi_level(func, **kwargs):
    """
    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, nbody_gufunc

    ptype = kwargs['ptype']
    return_wfn = kwargs.get('return_wfn', False)
    kwargs['return_wfn'] = True
    levels = {}
    for k, v in kwargs.pop('levels').items():
        if isinstance(k, str):
            levels[k.lower()] = v
        else:
            levels[k] = v
    supersystem = levels.pop('supersystem', False)
    molecule = kwargs.get('molecule', core.get_active_molecule())
    kwargs['bsse_type'] = [kwargs['bsse_type']] if isinstance(kwargs['bsse_type'], str) else kwargs['bsse_type']
    natoms = molecule.natom()

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

    if kwargs.get('charge_method', False) and not kwargs.get('embedding_charges', False):
        kwargs['embedding_charges'] = compute_charges(kwargs['charge_method'],
                                      kwargs.get('charge_type', 'MULLIKEN_CHARGES').upper(), molecule)

    for n in sorted(levels)[::-1]:
        molecule.set_name('%i' % n)
        kwargs_copy = kwargs.copy()
        kwargs_copy['max_nbody'] = n
        energy_bsse_dict = {b: 0 for b in kwargs['bsse_type']}
        if isinstance(levels[n], str):
            # If a new level of theory is provided, compute contribution
            ret, wfn = nbody_gufunc(func, levels[n], **kwargs_copy)
            wfns[n] = wfn
        else:
            # For the n-body contribution, use available data from the higher order levels[n]-body
            wfn = wfns[levels[n]]

        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 kwargs['bsse_type']:
                energy_bsse_dict[b] += sign * wfn.variable('%i%s' % (m, b.lower()))
            if ptype in ['gradient', 'hessian']:
                gradient_result += sign * np.array(wfn.variable('GRADIENT ' + str(m)))
                # Keep 1-body contribution to compute interaction data
                if n == 1:
                    gradient1 = np.array(wfn.variable('GRADIENT ' + str(m)))
            if ptype == 'hessian':
                hessian_result += sign * np.array(wfn.variable('HESSIAN ' + str(m)))
                if n == 1:
                    hessian1 = np.array(wfn.variable('HESSIAN ' + str(m)))
        energy_result += energy_bsse_dict[kwargs['bsse_type'][0]]
        for b in kwargs['bsse_type']:
            energy_body_contribution[b][n] = energy_bsse_dict[b]

    if supersystem:
        # Super system recovers higher order effects at a lower level
        molecule.set_name('supersystem')
        kwargs_copy = kwargs.copy()
        kwargs_copy.pop('bsse_type')
        kwargs_copy.pop('ptype')
        ret, wfn_super = func(supersystem, **kwargs_copy)
        core.clean()
        kwargs_copy = kwargs.copy()
        kwargs_copy['bsse_type'] = 'nocp'
        kwargs_copy['max_nbody'] = max(levels)
        # Subtract lower order effects to avoid double counting
        ret, wfn = nbody_gufunc(func, supersystem, **kwargs_copy)
        energy_result += wfn_super.energy() - wfn.variable(str(max(levels)))
        for b in kwargs['bsse_type']:
            energy_body_contribution[b][molecule.nfragments()] = wfn_super.energy() - wfn.variable(
                str(max(levels)))

        if ptype in ['gradient', 'hessian']:
            gradient_result += np.array(wfn_super.gradient()) - np.array(wfn.variable('GRADIENT ' + str(max(levels))))
        if ptype == 'hessian':
            hessian_result += np.array(wfn_super.hessian()) - np.array(wfn.variable('HESSIAN ' + str(max(levels))))
        levels['supersystem'] = supersystem

    for b in kwargs['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 = kwargs.get('embedding_charges', False) or kwargs.get('charge_method', False)
    for b in kwargs['bsse_type']:
        _print_nbody_energy(energy_body_dict[b], '%s-corrected multilevel many-body expansion' % b.upper(),
                            is_embedded)

    if not kwargs['return_total_data']:
        # Remove monomer contribution for interaction data
        energy_result -= energy_body_dict[kwargs['bsse_type'][0]][1]
        if ptype in ['gradient', 'hessian']:
            gradient_result -= gradient1
        if ptype == 'hessian':
            hessian_result -= hessian1
    wfn_out = core.Wavefunction.build(molecule, 'def2-svp')
    core.set_variable("CURRENT ENERGY", energy_result)
    wfn_out.set_variable("CURRENT ENERGY", energy_result)
    if gradient_result is not None:
        wfn_out.set_gradient(core.Matrix.from_array(gradient_result))
    if hessian_result is not None:
        wfn_out.set_hessian(core.Matrix.from_array(hessian_result) )
    ptype_result = eval(ptype + '_result')
    for b in kwargs['bsse_type']:
        for i in energy_body_dict[b]:
            wfn_out.set_variable(str(i) + b, energy_body_dict[b][i])

    if kwargs['return_wfn']:
        return (ptype_result, wfn_out)
    else:
        return ptype_result
Exemple #2
0
def multi_level(func, **kwargs):
    """
    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 nbody_gufunc
    from psi4.driver.driver_nbody import _print_nbody_energy

    ptype = kwargs['ptype']
    return_wfn = kwargs.get('return_wfn', False)
    kwargs['return_wfn'] = True
    levels = kwargs.pop('levels')
    for i in levels:
        if isinstance(i, str): levels[i.lower()] = levels.pop(i)
    supersystem = levels.pop('supersystem', False)
    molecule = kwargs.get('molecule', core.get_active_molecule())
    kwargs['bsse_type'] = [kwargs['bsse_type']] if isinstance(kwargs['bsse_type'], str) else kwargs['bsse_type']
    natoms = molecule.natom()

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

    if kwargs.get('charge_method', False) and not kwargs.get('embedding_charges', False):
        kwargs['embedding_charges'] = compute_charges(kwargs['charge_method'],
                                      kwargs.get('charge_type', 'MULLIKEN_CHARGES').upper(), molecule)

    for n in sorted(levels)[::-1]:
        molecule.set_name('%i' %n)
        kwargs_copy = kwargs.copy()
        kwargs_copy['max_nbody'] = n
        energy_bsse_dict = {b: 0 for b in kwargs['bsse_type']}
        if isinstance(levels[n], str):
            # If a new level of theory is provided, compute contribution
            ret, wfn = nbody_gufunc(func, levels[n], **kwargs_copy)
            wfns[n] = wfn
        else:
            # For the n-body contribution, use available data from the higher order levels[n]-body
            wfn = wfns[levels[n]]

        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 kwargs['bsse_type']:
                energy_bsse_dict[b] += sign * wfn.variable('%i%s' % (m, b.lower()))
            if ptype in ['gradient', 'hessian']:
                gradient_result += sign * np.array(wfn.variable('GRADIENT ' + str(m)))
                # Keep 1-body contribution to compute interaction data
                if n == 1:
                    gradient1 = np.array(wfn.variable('GRADIENT ' + str(m)))
            if ptype == 'hessian':
                hessian_result += sign * np.array(wfn.variable('HESSIAN ' + str(m)))
                if n == 1:
                    hessian1 = np.array(wfn.variable('HESSIAN ' + str(m)))
        energy_result += energy_bsse_dict[kwargs['bsse_type'][0]]
        for b in kwargs['bsse_type']:
            energy_body_contribution[b][n] = energy_bsse_dict[b]

    if supersystem:
        # Super system recovers higher order effects at a lower level
        molecule.set_name('supersystem')
        kwargs_copy = kwargs.copy()
        kwargs_copy.pop('bsse_type')
        kwargs_copy.pop('ptype')
        ret, wfn_super = func(supersystem, **kwargs_copy)
        core.clean()
        kwargs_copy = kwargs.copy()
        kwargs_copy['bsse_type'] = 'nocp'
        kwargs_copy['max_nbody'] = max(levels)
        # Subtract lower order effects to avoid double counting
        ret, wfn = nbody_gufunc(func, supersystem, **kwargs_copy)
        energy_result += wfn_super.energy() - wfn.variable(str(max(levels)))
        for b in kwargs['bsse_type']:
            energy_body_contribution[b][molecule.nfragments()] = wfn_super.energy() - wfn.variable(
                str(max(levels)))

        if ptype in ['gradient', 'hessian']:
            gradient_result += np.array(wfn_super.gradient()) - np.array(wfn.variable('GRADIENT ' + str(max(levels))))
        if ptype == 'hessian':
            hessian_result += np.array(wfn_super.hessian()) - np.array(wfn.variable('HESSIAN ' + str(max(levels))))
        levels['supersystem'] = supersystem

    for b in kwargs['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 = kwargs.get('embedding_charges', False) or kwargs.get('charge_method', False)
    for b in kwargs['bsse_type']:
        _print_nbody_energy(energy_body_dict[b], '%s-corrected multilevel many-body expansion' % b.upper(),
                            is_embedded)

    if not kwargs['return_total_data']:
        # Remove monomer cotribution for interaction data
        energy_result -= energy_body_dict[kwargs['bsse_type'][0]][1]
        if ptype in ['gradient', 'hessian']:
            gradient_result -= gradient1
        if ptype == 'hessian':
            hessian_result -= hessian1
    wfn_out = core.Wavefunction.build(molecule, 'def2-svp')
    core.set_variable("CURRENT ENERGY", energy_result)
    wfn_out.set_variable("CURRENT ENERGY", energy_result)
    gradient_result = core.Matrix.from_array(gradient_result) if gradient_result is not None else None
    wfn_out.set_gradient(gradient_result)
    hessian_result = core.Matrix.from_array(hessian_result) if hessian_result is not None else None
    wfn_out.set_hessian(hessian_result)
    ptype_result = eval(ptype + '_result')
    for b in kwargs['bsse_type']:
        for i in energy_body_dict[b]:
            wfn_out.set_variable(str(i) + b, energy_body_dict[b][i])

    if kwargs['return_wfn']:
        return (ptype_result, wfn_out)
    else:
        return ptype_result
def prepare_results(self,
                    client: Optional["FractalClient"] = None
                    ) -> Dict[str, Any]:
    """Use different levels of theory for different n-body levels.

    See ManyBodyComputer.prepare_results() for how this fits in.
    TODO: incorporate function into class.

    """
    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))

    # Get numerical label (index) for supersystem tasks
    sup_level = 0
    levels = []
    for n, i in enumerate(self.nbodies_per_mc_level):
        if 'supersystem' not in i:
            levels.append(int(n + 1))
        else:
            sup_level = n + 1

    nbody_list = self.nbodies_per_mc_level
    quiet = True

    for l in sorted(levels)[::-1]:
        self.quiet = quiet
        self.max_nbody = nbody_list[l - 1][-1]
        results = {
            k: v
            for k, v in self.task_list.items() if k.startswith(str(l))
        }
        results = self.prepare_results(results=results, client=client)

        for n in nbody_list[l - 1][::-1]:
            energy_bsse_dict = {b: 0 for b in self.bsse_type}

            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':
                    hessian_result += sign * results[f'{ptype}_body_dict'][m]
                    gradient_result += sign * results['gradient_body_dict'][m]
                    if n == 1:
                        hessian1 = results[f'{ptype}_body_dict'][n]
                        gradient1 = results['gradient_body_dict'][n]

                elif ptype == 'gradient':
                    gradient_result += sign * results[f'{ptype}_body_dict'][m]
                    # Keep 1-body contribution to compute interaction data
                    if n == 1:
                        gradient1 = results[f'{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.nfragments)).get_results(
                                                 client=client)
        self.max_nbody = max(levels)

        # Compute components at supersytem level of theory
        self.nbodies_per_mc_level.append(levels)
        component_result = {
            k: v
            for k, v in self.task_list.items() if k.startswith(str(sup_level))
        }
        components = self.prepare_results(results=component_result,
                                          client=client)

        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.extras.qcvars[
                'CURRENT GRADIENT'] - components['gradient_body_dict'][
                    self.max_nbody]
            hessian_result += supersystem_result.return_result - components[
                f'{ptype}_body_dict'][self.max_nbody]

        elif ptype == 'gradient':
            gradient_result += np.array(
                supersystem_result.return_result).reshape(
                    (-1, 3)) - components[f'{ptype}_body_dict'][self.max_nbody]

    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]
            ])

    is_embedded = self.embedding_charges
    for b in self.bsse_type:
        _print_nbody_energy(
            energy_body_dict[b],
            f"{b.upper()}-corrected multilevel many-body expansion",
            self.nfragments, is_embedded)

    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

    energy_body_dict = {
        str(k) + b: v
        for b in energy_body_dict for k, v in energy_body_dict[b].items()
    }

    nbody_results = {
        "ret_energy": energy_result,
        "ret_ptype": locals()[ptype + '_result'],
        "energy_body_dict": energy_body_dict,
    }
    return nbody_results