Skip to content

Commit

Permalink
added virial observable to toplevel api
Browse files Browse the repository at this point in the history
  • Loading branch information
@clonker
clonker committed Jan 23, 2018
1 parent 20bf3e3 commit 14a4ddc
Showing 1 changed file with 38 additions and 30 deletions.
68 changes: 38 additions & 30 deletions wrappers/python/src/python/readdy/api/registry/observables.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
@author: clonker
"""


def _parse_save_args(save_args):
assert save_args is None or isinstance(save_args, (dict, bool)), \
"save can only be None, bool, or a dictionary, not {}".format(type(save_args))
Expand All @@ -35,6 +36,7 @@ def _parse_save_args(save_args):
assert "name" in save_args.keys(), "save needs to have a \"name\" key"
return save_args["name"], save_args["chunk_size"]


class Observables(object):
def __init__(self, simulation):
self._simulation = simulation
Expand All @@ -56,23 +58,22 @@ def rdf(self, stride, bin_borders, types_count_from, types_count_to, particle_to
:param save: dictionary containing `name` and `chunk_size` or None to not save the observable to file
"""
if isinstance(save, str) and save == 'default':
save = {"name": "rdf", "chunk_size": 100}
save = {"name": "rdf", "chunk_size": 1000}
if isinstance(types_count_from, str):
types_count_from = [types_count_from]
if isinstance(types_count_to, str):
types_count_to = [types_count_to]
bin_borders = self._simulation._unit_conf.convert(bin_borders, self._simulation.length_unit)
particle_to_density = self._simulation._unit_conf.convert(particle_to_density,
1 / self._simulation.length_unit**3)
1 / self._simulation.length_unit ** 3)
assert all([isinstance(x, str) for x in types_count_from]), \
"types_count_from={} has an invalid type".format(types_count_from)
assert all([isinstance(x, str) for x in types_count_to]), \
"types_count_to={} has an invalid type".format(types_count_to)
handle = self._sim.register_observable_radial_distribution(stride, bin_borders, types_count_from, types_count_to,
handle = self._sim.register_observable_radial_distribution(stride, bin_borders, types_count_from,
types_count_to,
particle_to_density, callback)
name, chunk_size = _parse_save_args(save)
if name is not None and chunk_size is not None:
self._observable_handles.append((name, chunk_size, handle))
self._add_observable_handle(*_parse_save_args(save), handle)

def reactions(self, stride, callback=None, save='default'):
"""
Expand All @@ -83,12 +84,10 @@ def reactions(self, stride, callback=None, save='default'):
:param save: dictionary containing `name` and `chunk_size` or None to not save the observable to file
"""
if isinstance(save, str) and save == 'default':
save = {"name": "reactions", "chunk_size": 100}
save = {"name": "reactions", "chunk_size": 500}
handle = self._sim.register_observable_reactions(stride, callback)

name, chunk_size = _parse_save_args(save)
if name is not None and chunk_size is not None:
self._observable_handles.append((name, chunk_size, handle))
self._add_observable_handle(*_parse_save_args(save), handle)

def particle_positions(self, stride, types=None, callback=None, save='default'):
"""
Expand All @@ -106,9 +105,7 @@ def particle_positions(self, stride, types=None, callback=None, save='default'):
types = []
handle = self._sim.register_observable_particle_positions(stride, types, callback)

name, chunk_size = _parse_save_args(save)
if name is not None and chunk_size is not None:
self._observable_handles.append((name, chunk_size, handle))
self._add_observable_handle(*_parse_save_args(save), handle)

def particles(self, stride, callback=None, save='default'):
"""
Expand All @@ -123,9 +120,7 @@ def particles(self, stride, callback=None, save='default'):

handle = self._sim.register_observable_particles(stride, callback)

name, chunk_size = _parse_save_args(save)
if name is not None and chunk_size is not None:
self._observable_handles.append((name, chunk_size, handle))
self._add_observable_handle(*_parse_save_args(save), handle)

def number_of_particles(self, stride, types=None, callback=None, save='default'):
"""
Expand All @@ -143,9 +138,7 @@ def number_of_particles(self, stride, types=None, callback=None, save='default')
types = []
handle = self._sim.register_observable_n_particles(stride, types, callback)

name, chunk_size = _parse_save_args(save)
if name is not None and chunk_size is not None:
self._observable_handles.append((name, chunk_size, handle))
self._add_observable_handle(*_parse_save_args(save), handle)

def energy(self, stride, callback=None, save='default'):
"""
Expand All @@ -156,13 +149,11 @@ def energy(self, stride, callback=None, save='default'):
:param save: dictionary containing `name` and `chunk_size` or None to not save the observable to file
"""
if isinstance(save, str) and save == 'default':
save = {"name": "energy", "chunk_size": 100}
save = {"name": "energy", "chunk_size": 10000}

handle = self._sim.register_observable_energy(stride, callback)

name, chunk_size = _parse_save_args(save)
if name is not None and chunk_size is not None:
self._observable_handles.append((name, chunk_size, handle))
self._add_observable_handle(*_parse_save_args(save), handle)

def forces(self, stride, types=None, callback=None, save='default'):
"""
Expand All @@ -180,9 +171,7 @@ def forces(self, stride, types=None, callback=None, save='default'):
types = []
handle = self._sim.register_observable_forces(stride, types, callback)

name, chunk_size = _parse_save_args(save)
if name is not None and chunk_size is not None:
self._observable_handles.append((name, chunk_size, handle))
self._add_observable_handle(*_parse_save_args(save), handle)

def reaction_counts(self, stride, callback=None, save='default'):
"""
Expand All @@ -194,10 +183,29 @@ def reaction_counts(self, stride, callback=None, save='default'):
:param save: dictionary containing `name` and `chunk_size` or None to not save the observable to file
"""
if isinstance(save, str) and save == 'default':
save = {"name": "reaction_counts", "chunk_size": 100}
save = {"name": "reaction_counts", "chunk_size": 500}

handle = self._sim.register_observable_reaction_counts(stride, callback)

name, chunk_size = _parse_save_args(save)
if name is not None and chunk_size is not None:
self._observable_handles.append((name, chunk_size, handle))
self._add_observable_handle(*_parse_save_args(save), handle)

def virial(self, stride, callback=None, save='default'):
"""
Records the virial tensor of the system.
:param stride: skip `stride` time steps before evaluating the observable again
:param callback: callback function that takes the current virial as argument in terms of a (3,3) numpy array
:param save: dictionary containing `name` and `chunk_size` or None to not save the observable to file
"""
if isinstance(save, str) and save == 'default':
save = {"name": "virial", "chunk_size": 100}

handle = self._sim.register_observable_virial(stride, lambda x: callback(x.toarray()))
self._add_observable_handle(*_parse_save_args(save), handle)


def _add_observable_handle(self, save_name, save_chunk_size, handle):
if save_name is not None and save_chunk_size is not None:
if next((n for n, c, h in self._observable_handles if n == save_name), None) is not None:
raise RuntimeError("A observable with the name {} is already being recorded into the trajectory file."
.format(save_name))
self._observable_handles.append((save_name, save_chunk_size, handle))

0 comments on commit 14a4ddc

Please sign in to comment.