# -*- 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