Source code for pyretis.setup.createsystem

# -*- coding: utf-8 -*-
# Copyright (c) 2023, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""This module handles the set-up of initial positions and a box.

The initial positions can either be generated on a lattice, or it can
be read from a file.

Important methods defined here
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

set_up_box (:py:func:`.set_up_box`)
    Create a simulation box from simulation settings.

create_initial_positions (:py:func:`.create_initial_positions`)
    Get initial positions based on settings. This will either be
    read from a file or generated on a lattice.

create_system (:py:func:`.create_system`)
    Set up a system from given settings. This method will probably
    also need to set/get the initial positions and velocities for
    the particles and set up the simulation box.

create_velocities (:py:func:`.create_velocities`)
    Create velocities from settings for a system with particles.

initial_positions_file (:py:func:`.initial_positions_file`)
    Get initial positions from a file.

initial_positions_lattice (:py:func:`.initial_positions_lattice`)
    Get initial positions by generating a lattice.

"""
import logging
import os
import numpy as np
from pyretis.tools import generate_lattice
from pyretis.core.box import create_box
from pyretis.core.system import System
from pyretis.core.particles import Particles, get_particle_type
from pyretis.core.units import CONVERT
from pyretis.inout.settings import look_for_input_files
from pyretis.inout.formats.snapshot import read_txt_snapshots
from pyretis.inout.formats.xyz import read_xyz_file
from pyretis.inout.formats.gromacs import read_gromacs_file, read_gromos96_file
from pyretis.engines.openmm import make_pyretis_system as openmm_system
logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
logger.addHandler(logging.NullHandler())


__all__ = ['create_initial_positions', 'create_system', 'create_velocities',
           'set_up_box', 'initial_positions_file', 'initial_positions_lattice']


PERIODIC_TABLE = {'H': 1.007975, 'He': 4.002602, 'Li': 6.9675,
                  'Be': 9.0121831, 'B': 10.8135, 'C': 12.0106,
                  'N': 14.006855, 'O': 15.9994, 'F': 18.998403163,
                  'Ne': 20.1797, 'Na': 22.98976928, 'Mg': 24.3055,
                  'Al': 26.9815385, 'Si': 28.085, 'P': 30.973761998,
                  'S': 32.0675, 'Cl': 35.4515, 'Ar': 39.948,
                  'K': 39.0983, 'Ca': 40.078, 'Sc': 44.955908,
                  'Ti': 47.867, 'V': 50.9415, 'Cr': 51.9961,
                  'Mn': 54.938044, 'Fe': 55.845, 'Co': 58.933194,
                  'Ni': 58.6934, 'Cu': 63.546, 'Zn': 65.38,
                  'Ga': 69.723, 'Ge': 72.63, 'As': 74.921595,
                  'Se': 78.971, 'Br': 79.904, 'Kr': 83.798,
                  'Rb': 85.4678, 'Sr': 87.62, 'Y': 88.90584,
                  'Zr': 91.224, 'Nb': 92.90637, 'Mo': 95.95,
                  'Ru': 101.07, 'Rh': 102.9055, 'Pd': 106.42,
                  'Ag': 107.8682, 'Cd': 112.414, 'In': 114.818,
                  'Sn': 118.71, 'Sb': 121.76, 'Te': 127.6,
                  'I': 126.90447, 'Xe': 131.293, 'Cs': 132.90545196,
                  'Ba': 137.327, 'La': 138.90547, 'Ce': 140.116,
                  'Pr': 140.90766, 'Nd': 144.242, 'Sm': 150.36,
                  'Eu': 151.964, 'Gd': 157.25, 'Tb': 158.92535,
                  'Dy': 162.5, 'Ho': 164.93033, 'Er': 167.259,
                  'Tm': 168.93422, 'Yb': 173.045, 'Lu': 174.9668,
                  'Hf': 178.49, 'Ta': 180.94788, 'W': 183.84,
                  'Re': 186.207, 'Os': 190.23, 'Ir': 192.217,
                  'Pt': 195.084, 'Au': 196.966569, 'Hg': 200.592,
                  'Tl': 204.3835, 'Pb': 207.2, 'Bi': 208.9804,
                  'Th': 232.0377, 'Pa': 231.03588, 'U': 238.02891}


# Variable that defines some information for reading input files:
# TODO: Move this to the place where the file readers are defined?
READFILE = {'xyz': {'reader': read_xyz_file,
                    'units': {'length': 'A', 'velocity': 'A/fs'}},
            'txt': {'reader': read_txt_snapshots,
                    'units': None},
            'gro': {'reader': read_gromacs_file,
                    'units': {'length': 'nm', 'velocity': 'nm/ps'}},
            'g96': {'reader': read_gromos96_file,
                    'units': {'length': 'nm', 'velocity': 'nm/ps'}}}


def list_get(input_list, index):
    """Get an item from a list and handle out-of bounds errors.

    This method is intended to be used when we are picking items from
    a list and possibly we want a number of items which is larger than
    the number of items in the list. Here, we then just return the last
    element.

    Parameters
    ----------
    input_list : list
        The list to pick from.
    index : integer
        The index to pick.

    """
    try:
        return input_list[index]
    except IndexError:
        return input_list[-1]


[docs]def _assign_mass_from_file(filename, unit): """Assign masses and imasses to particles from a configuration file. This function is useful when velocities will be internally modified while using an external engine. Parameters ---------- filename : string The configuration file. unit : string The system of units. """ snapshot, _ = _get_snapshot_from_file({'input_file': filename}, unit) masses = [] for i, atomname in enumerate(snapshot['atomname']): masses.append([guess_particle_mass(i+1, atomname, unit)*1.]) return masses, np.reciprocal(masses)
def guess_particle_mass(particle_no, particle_type, unit): """Guess a particle mass from it's type. Parameters ---------- particle_no : integer Just used to identify the particle number. particle_type : string Used to identify the particle. unit : string The system of units. This is used in case we try to get the mass from the periodic table where the units are in `g/mol`. """ logger.info(('Mass not specified for particle no. %i\n' 'Will guess from particle type "%s"'), particle_no, particle_type) mass = PERIODIC_TABLE.get(particle_type, None) if mass is None: particle_mass = 1.0 logger.info(('-> Could not find mass. ' 'Assuming %f (internal units)'), particle_mass) else: particle_mass = CONVERT['mass']['g/mol', unit] * mass logger.info(('-> Using a mass of %f g/mol ' '(%f in internal units)'), mass, particle_mass) return particle_mass
[docs]def initial_positions_lattice(settings): """Generate initial positions based on given settings. We assume here the input values are given with the correct units as dictated by ``settings['system']['units']``. Parameters ---------- settings : dict The input settings for the simulation. Returns ------- particles : object like :py:class:`.Particles` The particles we created. size : list of floats A size for the region we created. This can be used to create a box. """ pos_settings = settings['particles']['position'] ptype = settings['particles'].get('ptype', [0]) pname = settings['particles'].get('name', ['Ar']) lattice, size = generate_lattice( pos_settings['generate'].lower(), pos_settings.get('repeat'), lcon=pos_settings.get('lcon', None), density=pos_settings.get('density', None) ) ndim = settings['system'].get('dimensions', None) if ndim is None: ndim = len(size) else: if ndim != len(size): msgtxt = ('Inconsistent dimenaions in settings and generated ' f'lattice.\nSettings gives {ndim}D while the generated ' f'lattice is {len(size)}D.') logger.error(msgtxt) raise ValueError(msgtxt) particles = Particles(dim=ndim) for i, pos in enumerate(lattice): particle_type = list_get(ptype, i) particle_name = list_get(pname, i) # Infer the mass from the input masses, or try to get it # from the periodic table: p_mass = settings['particles'].get('mass', {}).get(particle_name) if p_mass is None: p_mass = guess_particle_mass(i + 1, particle_name, settings['system']['units']) particles.add_particle( pos, np.zeros_like(pos), np.zeros_like(pos), mass=p_mass, name=particle_name, ptype=particle_type ) logger.info('Generated %i particles on lattice "%s".', particles.npart, pos_settings['generate'].lower()) logger.info('Lattice is %iD.', ndim) size = np.array(size) box = {'low': size[:, 0], 'high': size[:, 1]} return particles, box
[docs]def _get_snapshot_from_file(pos_settings, units): """Get a configuration snapshot from a file. This snapshot will be used to set up the initial configuration. Parameters ---------- pos_settings : dict A dict with information on what we should read. units : string The internal units. Returns ------- snapshot : dict The snapshot we found in the file. It will at least have the keys with the positions ('x', 'y', 'z') and atom name 'atomnames'. It may have information about velocities ('vx', 'vy', 'vz') and the box ('box'). convert : dict Dictionary with conversion factors to internal units. """ filename = pos_settings.get('input_file', None) if filename is None: msg = ('Requested reading (initial) configuration from file, ' 'but no input file given!') logger.error(msg) raise ValueError(msg) fmt = pos_settings.get('format', os.path.splitext(filename)[1][1:]) snaps = [] convert = None if fmt not in READFILE: msg = (f'Input configuration "{filename}" has unknown ' f'format "{fmt}".') logger.error(msg) logger.error('Supported formats are: %s.', list(READFILE)) raise ValueError(msg) reader = READFILE[fmt]['reader'] read_units = READFILE[fmt]['units'] if read_units is None: convert = {'length': 1.0, 'velocity': 1.0} else: convert = { 'length': CONVERT['length'][read_units['length'], units], 'velocity': CONVERT['velocity'][read_units['velocity'], units] } logger.info( 'Reading initial configuration from "%s" (format: "%s").', filename, fmt, ) snaps = [snap for snap in reader(filename)] snapshot = None if len(snaps) == 1: snapshot = snaps[0] elif len(snaps) > 1: msg = ('Found several frames ({}) in input file.' ' Will use the last one!').format(len(snaps)) logger.warning(msg) snapshot = snaps[-1] else: msg = ('Could not find any configurations in ' 'input file: {}').format(filename) logger.error(msg) raise ValueError(msg) return snapshot, convert
[docs]def initial_positions_file(settings): """Get initial positions from an input file. Parameters ---------- settings : dict The input settings for the simulation. Returns ------- particles : object like :py:class:`.Particles` The particles we created. size : list of floats A size for the region we created. This can be used to create a box. vel_read : boolean True if we read velocities from the input file. """ ndim = settings['system'].get('dimensions', 3) pos_settings = settings['particles']['position'] ptype = settings['particles'].get('ptype', None) pname = settings['particles'].get('name', None) pmass = settings['particles'].get('mass', {}) ptypes = {} # To automatically set particle types based on name. snapshot, convert = _get_snapshot_from_file(pos_settings, settings['system']['units']) vel_read = False particles = Particles(dim=ndim) for i, atomname in enumerate(snapshot['atomname']): pos = [] vel = [] for key in ['x', 'y', 'z'][:ndim]: pos.append(snapshot[key][i]) vel_key = 'v{}'.format(key) if vel_key in snapshot: vel.append(snapshot[vel_key][i]) pos = np.array(pos) * convert['length'] if len(vel) != ndim: vel = np.zeros_like(pos) else: vel = np.array(vel) * convert['velocity'] vel_read = True # Get particle type from the atom names or from input list: if ptype is None: if atomname not in ptypes: ptypes[atomname] = len(ptypes) particle_type = ptypes[atomname] else: particle_type = list_get(ptype, i) if pname is None: particle_name = atomname else: particle_name = list_get(pname, i) # Infer the mass from the input masses, or try to get it # from the periodic table: try: particle_mass = pmass[particle_name] except KeyError: particle_mass = guess_particle_mass(i + 1, particle_name, settings['system']['units']) particles.add_particle(pos, vel, np.zeros_like(pos), mass=particle_mass, name=particle_name, ptype=particle_type) try: box = {'cell': [i * convert['length'] for i in snapshot['box']]} if ndim < 3: box['cell'] = box['cell'][:ndim] except (KeyError, IndexError, TypeError) as err: logger.debug('No box read from file: %s.', err) box = None logger.info('Read %d particle(s) from "%s".', particles.npart, pos_settings['input_file']) if vel_read: logger.info('Read velocities from file: "%s".', pos_settings['input_file']) return particles, box, vel_read
[docs]def create_initial_positions(settings): """Set up the initial positions from the given settings. The settings can specify the initial positions as a file or to be generated on a lattice by PyRETIS. Parameters ---------- settings : dict Settings for creating the initial positions. Returns ------- out[0] : object like :py:class:`.Particles` The particles we created. out[1] : list The size associated with the particles. Can be used to create a box. out[2] : boolean True if we have read/created velocities different from just zeros. This is only True if we have read from a file with velocities. """ logger.debug('Settings used for initial positions: %s', settings['particles']['position']) particles = None if 'generate' in settings['particles']['position']: particles, box = initial_positions_lattice(settings) return particles, box, False if 'input_file' in settings['particles']['position']: # First check if we need to add a path to the file: filename = settings['particles']['position']['input_file'] if (not os.path.isfile(filename) and 'exe_path' in settings['simulation']): filename = os.path.join(settings['simulation']['exe_path'], filename) settings['particles']['position']['input_file'] = filename particles, box, vel = initial_positions_file(settings) return particles, box, vel msg = 'Unknown settings for initial positions: {}' msgtxt = msg.format(settings['particles']['position']) logger.error(msgtxt) raise ValueError(msgtxt)
[docs]def set_up_box(settings, boxs, dim=3): """Set up a box from given settings. Parameters ---------- settings : dict The dict with the simulation settings. boxs : dict or None If no box settings are given, we can still create a box, inferred from the positions of the particles. This dict contains the settings to do so. dim : integer, optional Number of dimensions for the box. This is used only as a last resort when no information about the box is given. Returns ------- box : object like :py:class:`.BoxBase` or None The box if we managed to create it, otherwise None. """ msg = 'Box created {}:\n{}' box = None if settings.get('box', None) is not None: box = create_box(**settings['box']) msgtxt = msg.format('from settings', box) logger.info(msgtxt) debugtxt = 'Settings used:\n{}'.format(settings['box']) logger.debug(debugtxt) else: if boxs is not None: box = create_box(**boxs) msgtxt = msg.format('from initial positions', box) logger.info(msgtxt) msgwarn = 'The box was assumed periodic in all directions.' logger.warning(msgwarn) else: if dim > 0: box = create_box(periodic=[False]*dim) msgtxt = msg.format('without specifications', box) logger.info(msgtxt) msgwarn = 'The box was assumed nonperiodic in all directions.' logger.warning(msgwarn) return box
[docs]def create_velocities(system, settings, vel): """Create velocities from settings for a system. Parameters ---------- system : object like :py:class:`.System` The system to create velocities for. It's needed since we need to know the degrees of freedom. settings : dict Settings to use for creating the velocities. vel : boolean If True, we already read velocities. They will now be overwritten. We just make some warnings about this. Returns ------- out : boolean True if we actually generated velocities. """ vel_settings = settings['particles'].get('velocity', {}) if vel: logger.info( 'Velocities read from input configuration (temperature: %6.2g).', system.calculate_temperature(), ) if 'generate' in vel_settings: if vel: logger.warning( 'Will generate and overwrite velocities already set.' ) gen_settings = {'distribution': vel_settings['generate']} for key in ('seed', 'momentum', 'temperature', 'rgen'): try: gen_settings[key] = vel_settings[key] except KeyError: pass system.generate_velocities(**gen_settings) logger.info( 'Generated new velocities with average temperature: %6.2g', system.calculate_temperature(), ) logger.debug('Settings used for generating velocities: %s', gen_settings) return True if 'scale' in vel_settings: target = vel_settings['scale'] # Just set the velocities to some temperature for now. # The scaling is done later by calling `system.extra_setup()`. gen_settings = { 'distribution': 'maxwell', 'momentum': system.particles.npart != 1, 'temperature': settings['system'].get('temperature', 1.0), } system.generate_velocities(**gen_settings) msg = 'Scaling velocities to total energy {}'.format(target) logger.debug(msg) system.post_setup.append(('rescale_velocities', (target,))) return True return False
[docs]def create_system(settings): """Create a system from input settings. Parameters ---------- settings : dict The dict with the simulation settings. Returns ------- system : object like :py:class:`.System` The system object we create here. """ if 'restart' in settings: system = System() system.load_restart_info(settings['restart']['system']) return system vel = None if ('engine' not in settings or settings['engine'].get('type', 'internal') == 'internal'): particles, box, vel = create_initial_positions(settings) box = set_up_box(settings, box, dim=particles.dim) elif settings['engine'].get('type', None).lower() == 'openmm': return openmm_system(settings) else: # Engine is not None and not internal => external. klass = get_particle_type('external') particles = klass(dim=3) if 'input_path' in settings['engine']: required_file = {} # Get the engine input files if 'cp2k' in settings['engine']['class'].lower(): required_file = {'conf': 'initial.{}'.format( settings['engine'].get('cp2k_format', 'xyz'))} elif 'gromacs' in settings['engine']['class'].lower(): required_file = { 'conf': f"conf.{settings['engine'].get('gmx_format', 'gro')}"} input_file = look_for_input_files(settings['engine']['input_path'], required_file) particles.set_pos((input_file['conf'], None)) # gromacs tests fail without this if-statement. # _get_snapshot_from_file() has issues with reading .g96 format. if settings['engine']['class'].lower() == 'cp2k': particles.mass, particles.imass = _assign_mass_from_file( input_file['conf'], settings['system']['units']) else: particles.set_pos((None, None)) particles.set_vel(False) box = None system = System( temperature=settings['system']['temperature'], units=settings['system']['units'], box=box ) system.particles = particles # figure out what to do with velocities: if 'position' in settings['particles'] and \ 'input_file' not in settings['particles']['position']: vel_gen = create_velocities(system, settings, vel) if not (vel_gen or vel): logger.warning('Velocities not created/read: Set to zero!') return system