How to use the schnetpack.Properties function in schnetpack

To help you get started, we’ve selected a few schnetpack examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github atomistic-machine-learning / schnetpack / src / schnetpack / data / atoms.py View on Github external
reformated into SchNetPack input format.
    """
    if output is None:
        inputs = {}
    else:
        inputs = output

    # Elemental composition
    cell = np.array(atoms.cell.array, dtype=np.float32)  # get cell array

    inputs[Properties.Z] = torch.LongTensor(atoms.numbers.astype(np.int))
    positions = atoms.positions.astype(np.float32)
    if center_positions:
        positions -= atoms.get_center_of_mass()
    inputs[Properties.R] = torch.FloatTensor(positions)
    inputs[Properties.cell] = torch.FloatTensor(cell)

    # get atom environment
    nbh_idx, offsets = environment_provider.get_environment(atoms)

    # Get neighbors and neighbor mask
    inputs[Properties.neighbors] = torch.LongTensor(nbh_idx.astype(np.int))

    # Get cells
    inputs[Properties.cell] = torch.FloatTensor(cell)
    inputs[Properties.cell_offset] = torch.FloatTensor(offsets.astype(np.float32))

    # If requested get neighbor lists for triples
    if collect_triples:
        nbh_idx_j, nbh_idx_k, offset_idx_j, offset_idx_k = collect_atom_triples(nbh_idx)
        inputs[Properties.neighbor_pairs_j] = torch.LongTensor(nbh_idx_j.astype(np.int))
        inputs[Properties.neighbor_pairs_k] = torch.LongTensor(nbh_idx_k.astype(np.int))
github atomistic-machine-learning / schnetpack / src / schnetpack / md / calculators / sgdml_calculator.py View on Github external
def __init__(
        self,
        model,
        required_properties=[Properties.energy, Properties.forces],
        force_handle=Properties.forces,
        position_conversion=1.0 / MDUnits.angs2bohr,
        force_conversion=1.0 / MDUnits.Ha2kcalpmol / MDUnits.angs2bohr,
        property_conversion={},
        detach=True,
    ):
        super(SGDMLCalculator, self).__init__(
            required_properties,
            force_handle,
            position_conversion,
            force_conversion,
            property_conversion,
            detach,
        )

        self.model = model
github atomistic-machine-learning / schnetpack / src / schnetpack / md / calculators / orca_calculator.py View on Github external
from schnetpack.md.utils import MDUnits
from schnetpack.md.parsers.orca_parser import OrcaMainFileParser

from schnetpack import Properties

import torch

import os
import subprocess

from ase.data import chemical_symbols
from ase import Atoms


class OrcaCalculator(QMCalculator):
    is_atomistic = [Properties.forces, Properties.shielding]

    def __init__(
        self,
        required_properties,
        force_handle,
        compdir,
        qm_executable,
        orca_template,
        orca_parser=OrcaMainFileParser,
        position_conversion=1.0 / MDUnits.angs2bohr,
        force_conversion=1.0,
        property_conversion={},
        queuer=None,
        adaptive=False,
        basename="input",
    ):
github atomistic-machine-learning / schnetpack / src / schnetpack / md / parsers / orca_parser.py View on Github external
return self.formatters.format(self.parsed)

    def reset(self):
        """
        Reset state of parser
        """
        self.read = False
        self.parsed = None


class OrcaMainFileParser(OrcaOutputParser):
    properties = [
        "atoms",
        Properties.forces,
        Properties.energy,
        Properties.dipole_moment,
        Properties.polarizability,
        Properties.shielding,
    ]

    starts = {
        "atoms": "CARTESIAN COORDINATES (ANGSTROEM)",
        Properties.forces: "CARTESIAN GRADIENT",
        Properties.energy: "FINAL SINGLE POINT ENERGY",
        Properties.dipole_moment: "Total Dipole Moment",
        Properties.polarizability: "The raw cartesian tensor (atomic units):",
        Properties.shielding: "CHEMICAL SHIFTS",
    }

    stops = {
        "atoms": "CARTESIAN COORDINATES (A.U.)",
        Properties.forces: "Difference to translation invariance",
github atomistic-machine-learning / G-SchNet / utility_functions.py View on Github external
for i in index_list:
            _i = i-n_tokens  # atom index if ignoring tokens
            while True:
                partial_mol = mol.copy()
                # get the type of the next atom
                next_type = pred_types[_i+j]
                # don't consider distance predictions at stop token prediction steps
                if next_type == stop_token:
                    dist_mask = torch.zeros(i).float()
                else:
                    dist_mask = torch.ones(i).float()
                # always consider the type prediction
                type_mask = torch.ones(i).float()
                # assemble neighborhood mask and neighbor indices
                neighbor_mask = torch.ones(i, i-1)
                neighbors = partial_mol[Properties.neighbors][:i, :i-1]
                # set position and labels of the current (focus) token (first atom)
                cur = focus[_i + j]
                pos = partial_mol[Properties.R][:i]
                label_idx = i if i < n_atoms else 0
                labels = partial_mol['_labels'][label_idx, :i]
                pos = torch.cat((pos[cur:cur+1], pos[1:]), 0)
                labels = torch.cat((labels[cur:cur+1], labels[1:]), 0)

                partial_mol.update(
                    {Properties.R: pos,
                     Properties.Z: partial_mol[Properties.Z][:i],
                     '_labels': labels,
                     '_dist_mask': dist_mask,
                     '_type_mask': type_mask,
                     Properties.neighbor_mask: neighbor_mask,
                     Properties.neighbors: neighbors,
github atomistic-machine-learning / G-SchNet / utility_functions.py View on Github external
def build_batch(i):
        amount = torch.sum(unfinished)  # only get predictions for unfinished molecules
        # build neighborhood and neighborhood mask
        neighbors_i = neighbors[:i, :i-1].expand(amount, -1, -1).contiguous()
        neighbor_mask = torch.ones_like(neighbors_i).float()
        # set position of focus token (first entry of positions)
        positions[unfinished, 0] = positions[unfinished, current_atoms[unfinished]]
        # center positions on currently focused atom (for localized grid)
        positions[unfinished, :i] -= \
            positions[unfinished, current_atoms[unfinished]][:, None, :]

        # build batch with data of the partial molecules
        batch = {
            Properties.R: positions[unfinished, :i],
            Properties.Z: atom_numbers[unfinished, :i],
            Properties.atom_mask: torch.zeros(amount, i, dtype=torch.float),
            Properties.neighbors: neighbors_i,
            Properties.neighbor_mask: neighbor_mask,
            Properties.cell_offset: torch.zeros(amount, i, max(i-1, 1), n_dims),
            Properties.cell: torch.zeros(amount, n_dims, n_dims),
            '_next_types': atom_numbers[unfinished, i],
            '_all_types': all_types.view(1, -1),
            '_type_mask': torch.ones(amount, i, dtype=torch.float),
        }

        # put batch into torch variables and on gpu
        batch = {
            k: v.to(device)
            for k, v in batch.items()
        }
        return batch
github atomistic-machine-learning / schnetpack / src / schnetpack / representation / schnet.py View on Github external
"""Compute atomic representations/embeddings.

        Args:
            inputs (dict of torch.Tensor): SchNetPack dictionary of input tensors.

        Returns:
            torch.Tensor: atom-wise representation.
            list of torch.Tensor: intermediate atom-wise representations, if
            return_intermediate=True was used.

        """
        # get tensors from input dictionary
        atomic_numbers = inputs[Properties.Z]
        positions = inputs[Properties.R]
        cell = inputs[Properties.cell]
        cell_offset = inputs[Properties.cell_offset]
        neighbors = inputs[Properties.neighbors]
        neighbor_mask = inputs[Properties.neighbor_mask]
        atom_mask = inputs[Properties.atom_mask]

        # get atom embeddings for the input atomic numbers
        x = self.embedding(atomic_numbers)

        if False and self.charged_systems and Properties.charge in inputs.keys():
            n_atoms = torch.sum(atom_mask, dim=1, keepdim=True)
            charge = inputs[Properties.charge] / n_atoms  # B
            charge = charge[:, None] * self.charge  # B x F
            x = x + charge

        # compute interatomic distance of every atom to its neighbors
        r_ij = self.distances(
            positions, neighbors, cell, cell_offset, neighbor_mask=neighbor_mask
github atomistic-machine-learning / G-SchNet / filter_generated.py View on Github external
idcs = np.where(d['valid'])[0]
            if len(idcs) < 1:
                i = end
                continue
            # shrink stats
            idx_id = stat_heads.index('id')
            idx_known = stat_heads.index('known')
            new_stats = stats[:, start:end]
            if 'new' in args.store and args.model_path is not None:
                idcs = idcs[np.where(new_stats[idx_known, idcs] == 0)[0]]
            new_stats = new_stats[:, idcs]
            new_stats[idx_id] = np.arange(len(new_stats[idx_id]))  # adjust ids
            shrunk_stats = np.hstack((shrunk_stats, new_stats))
            # shrink positions and atomic numbers
            shrunk_res[key] = {Properties.R: d[Properties.R][idcs],
                               Properties.Z: d[Properties.Z][idcs]}
            # store connectivity matrices if desired
            if 'connectivity' in args.store:
                shrunk_res[key].update(
                    {'connectivity': [d['connectivity'][k] for k in idcs]})
            i = end

        shrunk_res['stats'] = shrunk_stats
        res = shrunk_res

    # store results in new database
    # get filename that is not yet taken for db
    if os.path.isfile(target_db):
        file_name, _ = os.path.splitext(target_db)
        expand = 0
        while True:
            expand += 1