<commit_before># -*- coding: utf-8 -*-
# Run with one of these commands:
#    > OPENERP_ADDONS_PATH='../../addons/trunk' OPENERP_PORT=8069 \
#      OPENERP_DATABASE=yy PYTHONPATH=. python tests/test_ir_sequence.py
#    > OPENERP_ADDONS_PATH='../../addons/trunk' OPENERP_PORT=8069 \
#      OPENERP_DATABASE=yy nosetests tests/test_ir_sequence.py
#    > OPENERP_ADDONS_PATH='../../../addons/trunk' OPENERP_PORT=8069 \
#      OPENERP_DATABASE=yy PYTHONPATH=../:. unit2 test_ir_sequence
# This assume an existing database.
import psycopg2
import unittest2

import openerp
import common

DB = common.DB
ADMIN_USER_ID = common.ADMIN_USER_ID

def registry(model):
    return openerp.modules.registry.RegistryManager.get(DB)[model]

def cursor():
    return openerp.modules.registry.RegistryManager.get(DB).db.cursor()

class test_ir_sequence_standard(unittest2.TestCase):
    """ A few tests for a 'Standard' (i.e. PostgreSQL) sequence. """

    def test_ir_sequence_create(self):
        """ Try to create a sequence object. """
        cr = cursor()
        d = dict(code='test_sequence_type', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type', name='Test sequence')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        cr.commit()
        cr.close()

    def test_ir_sequence_search(self):
        """ Try a search. """
        cr = cursor()
        ids = registry('ir.sequence').search(cr, ADMIN_USER_ID, [], {})
        assert ids
        cr.commit()
        cr.close()

    def test_ir_sequence_draw(self):
        """ Try to draw a number. """
        cr = cursor()
        n = registry('ir.sequence').get(cr, ADMIN_USER_ID, 'test_sequence_type', {})
        assert n
        cr.commit()
        cr.close()

    def test_ir_sequence_draw_twice(self):
        """ Try to draw a number from two transactions. """
        cr0 = cursor()
        cr1 = cursor()
        n0 = registry('ir.sequence').get(cr0, ADMIN_USER_ID, 'test_sequence_type', {})
        assert n0
        n1 = registry('ir.sequence').get(cr1, ADMIN_USER_ID, 'test_sequence_type', {})
        assert n1
        cr0.commit()
        cr1.commit()
        cr0.close()
        cr1.close()

class test_ir_sequence_no_gap(unittest2.TestCase):
    """ Copy of the previous tests for a 'No gap' sequence. """

    def test_ir_sequence_create_no_gap(self):
        """ Try to create a sequence object. """
        cr = cursor()
        d = dict(code='test_sequence_type_2', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_2', name='Test sequence',
            implementation='no_gap')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        cr.commit()
        cr.close()

    def test_ir_sequence_draw_no_gap(self):
        """ Try to draw a number. """
        cr = cursor()
        n = registry('ir.sequence').get(cr, ADMIN_USER_ID, 'test_sequence_type_2', {})
        assert n
        cr.commit()
        cr.close()

    def test_ir_sequence_draw_twice_no_gap(self):
        """ Try to draw a number from two transactions.
        This is expected to not work.
        """
        cr0 = cursor()
        cr1 = cursor()
        cr1._default_log_exceptions = False # Prevent logging a traceback
        msg_re = '^could not obtain lock on row in relation "ir_sequence"$'
        with self.assertRaisesRegexp(psycopg2.OperationalError, msg_re):
            n0 = registry('ir.sequence').get(cr0, ADMIN_USER_ID, 'test_sequence_type_2', {})
            assert n0
            n1 = registry('ir.sequence').get(cr1, ADMIN_USER_ID, 'test_sequence_type_2', {})
        cr0.close()
        cr1.close()

class test_ir_sequence_change_implementation(unittest2.TestCase):
    """ Create sequence objects and change their ``implementation`` field. """

    def test_ir_sequence_1_create(self):
        """ Try to create a sequence object. """
        cr = cursor()
        d = dict(code='test_sequence_type_3', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_3', name='Test sequence', 'implementation': 'standard')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_4', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_4', name='Test sequence',
            implementation='no_gap')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        cr.commit()
        cr.close()

    def test_ir_sequence_2_write(self):
        cr = cursor()
        ids = registry('ir.sequence').search(cr, ADMIN_USER_ID,
            [('code', 'in', ['test_sequence_type_3', 'test_sequence_type_4'])], {})
        registry('ir.sequence').write(cr, ADMIN_USER_ID, ids,
            {'implementation': 'standard'}, {})
        registry('ir.sequence').write(cr, ADMIN_USER_ID, ids,
            {'implementation': 'no_gap'}, {})
        cr.commit()
        cr.close()

    def test_ir_sequence_3_unlink(self):
        cr = cursor()
        ids = registry('ir.sequence').search(cr, ADMIN_USER_ID,
            [('code', 'in', ['test_sequence_type_3', 'test_sequence_type_4'])], {})
        registry('ir.sequence').unlink(cr, ADMIN_USER_ID, ids, {})
        cr.commit()
        cr.close()

class test_ir_sequence_generate(unittest2.TestCase):
    """ Create sequence objects and generate some values. """

    def test_ir_sequence_create(self):
        """ Try to create a sequence object. """
        cr = cursor()
        d = dict(code='test_sequence_type_5', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_5', name='Test sequence')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        cr.commit()
        cr.close()

        cr = cursor()
        f = lambda *a: registry('ir.sequence').get(cr, ADMIN_USER_ID, 'test_sequence_type_5', {})
        assert all(str(x) == f() for x in xrange(1,1000))
        cr.commit()
        cr.close()

    def test_ir_sequence_create_no_gap(self):
        """ Try to create a sequence object. """
        cr = cursor()
        d = dict(code='test_sequence_type_6', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_6', name='Test sequence')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        cr.commit()
        cr.close()

        cr = cursor()
        f = lambda *a: registry('ir.sequence').get(cr, ADMIN_USER_ID, 'test_sequence_type_6', {})
        assert all(str(x) == f() for x in xrange(1,1000))
        cr.commit()
        cr.close()

if __name__ == '__main__':
    unittest2.main()

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<commit_msg>[MERGE] Merge fix of sequecne<commit_after># -*- coding: utf-8 -*-
# Run with one of these commands:
#    > OPENERP_ADDONS_PATH='../../addons/trunk' OPENERP_PORT=8069 \
#      OPENERP_DATABASE=yy PYTHONPATH=. python tests/test_ir_sequence.py
#    > OPENERP_ADDONS_PATH='../../addons/trunk' OPENERP_PORT=8069 \
#      OPENERP_DATABASE=yy nosetests tests/test_ir_sequence.py
#    > OPENERP_ADDONS_PATH='../../../addons/trunk' OPENERP_PORT=8069 \
#      OPENERP_DATABASE=yy PYTHONPATH=../:. unit2 test_ir_sequence
# This assume an existing database.
import psycopg2
import unittest2

import openerp
import common

DB = common.DB
ADMIN_USER_ID = common.ADMIN_USER_ID

def registry(model):
    return openerp.modules.registry.RegistryManager.get(DB)[model]

def cursor():
    return openerp.modules.registry.RegistryManager.get(DB).db.cursor()

class test_ir_sequence_standard(unittest2.TestCase):
    """ A few tests for a 'Standard' (i.e. PostgreSQL) sequence. """

    def test_ir_sequence_create(self):
        """ Try to create a sequence object. """
        cr = cursor()
        d = dict(code='test_sequence_type', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type', name='Test sequence')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        cr.commit()
        cr.close()

    def test_ir_sequence_search(self):
        """ Try a search. """
        cr = cursor()
        ids = registry('ir.sequence').search(cr, ADMIN_USER_ID, [], {})
        assert ids
        cr.commit()
        cr.close()

    def test_ir_sequence_draw(self):
        """ Try to draw a number. """
        cr = cursor()
        n = registry('ir.sequence').get(cr, ADMIN_USER_ID, 'test_sequence_type', {})
        assert n
        cr.commit()
        cr.close()

    def test_ir_sequence_draw_twice(self):
        """ Try to draw a number from two transactions. """
        cr0 = cursor()
        cr1 = cursor()
        n0 = registry('ir.sequence').get(cr0, ADMIN_USER_ID, 'test_sequence_type', {})
        assert n0
        n1 = registry('ir.sequence').get(cr1, ADMIN_USER_ID, 'test_sequence_type', {})
        assert n1
        cr0.commit()
        cr1.commit()
        cr0.close()
        cr1.close()

class test_ir_sequence_no_gap(unittest2.TestCase):
    """ Copy of the previous tests for a 'No gap' sequence. """

    def test_ir_sequence_create_no_gap(self):
        """ Try to create a sequence object. """
        cr = cursor()
        d = dict(code='test_sequence_type_2', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_2', name='Test sequence',
            implementation='no_gap')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        cr.commit()
        cr.close()

    def test_ir_sequence_draw_no_gap(self):
        """ Try to draw a number. """
        cr = cursor()
        n = registry('ir.sequence').get(cr, ADMIN_USER_ID, 'test_sequence_type_2', {})
        assert n
        cr.commit()
        cr.close()

    def test_ir_sequence_draw_twice_no_gap(self):
        """ Try to draw a number from two transactions.
        This is expected to not work.
        """
        cr0 = cursor()
        cr1 = cursor()
        cr1._default_log_exceptions = False # Prevent logging a traceback
        msg_re = '^could not obtain lock on row in relation "ir_sequence"$'
        with self.assertRaisesRegexp(psycopg2.OperationalError, msg_re):
            n0 = registry('ir.sequence').get(cr0, ADMIN_USER_ID, 'test_sequence_type_2', {})
            assert n0
            n1 = registry('ir.sequence').get(cr1, ADMIN_USER_ID, 'test_sequence_type_2', {})
        cr0.close()
        cr1.close()

class test_ir_sequence_change_implementation(unittest2.TestCase):
    """ Create sequence objects and change their ``implementation`` field. """

    def test_ir_sequence_1_create(self):
        """ Try to create a sequence object. """
        cr = cursor()
        d = dict(code='test_sequence_type_3', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_3', name='Test sequence')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_4', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_4', name='Test sequence',
            implementation='no_gap')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        cr.commit()
        cr.close()

    def test_ir_sequence_2_write(self):
        cr = cursor()
        ids = registry('ir.sequence').search(cr, ADMIN_USER_ID,
            [('code', 'in', ['test_sequence_type_3', 'test_sequence_type_4'])], {})
        registry('ir.sequence').write(cr, ADMIN_USER_ID, ids,
            {'implementation': 'standard'}, {})
        registry('ir.sequence').write(cr, ADMIN_USER_ID, ids,
            {'implementation': 'no_gap'}, {})
        cr.commit()
        cr.close()

    def test_ir_sequence_3_unlink(self):
        cr = cursor()
        ids = registry('ir.sequence').search(cr, ADMIN_USER_ID,
            [('code', 'in', ['test_sequence_type_3', 'test_sequence_type_4'])], {})
        registry('ir.sequence').unlink(cr, ADMIN_USER_ID, ids, {})
        cr.commit()
        cr.close()

class test_ir_sequence_generate(unittest2.TestCase):
    """ Create sequence objects and generate some values. """

    def test_ir_sequence_create(self):
        """ Try to create a sequence object. """
        cr = cursor()
        d = dict(code='test_sequence_type_5', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_5', name='Test sequence')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        cr.commit()
        cr.close()

        cr = cursor()
        f = lambda *a: registry('ir.sequence').get(cr, ADMIN_USER_ID, 'test_sequence_type_5', {})
        assert all(str(x) == f() for x in xrange(1,1000))
        cr.commit()
        cr.close()

    def test_ir_sequence_create_no_gap(self):
        """ Try to create a sequence object. """
        cr = cursor()
        d = dict(code='test_sequence_type_6', name='Test sequence type')
        c = registry('ir.sequence.type').create(cr, ADMIN_USER_ID, d, {})
        assert c
        d = dict(code='test_sequence_type_6', name='Test sequence')
        c = registry('ir.sequence').create(cr, ADMIN_USER_ID, d, {})
        assert c
        cr.commit()
        cr.close()

        cr = cursor()
        f = lambda *a: registry('ir.sequence').get(cr, ADMIN_USER_ID, 'test_sequence_type_6', {})
        assert all(str(x) == f() for x in xrange(1,1000))
        cr.commit()
        cr.close()

if __name__ == '__main__':
    unittest2.main()

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<|endoftext|>


<commit_before>import os
import numpy as np
from openmmtools.integrators import VVVRIntegrator

import simtk.unit as u
import simtk.openmm as openmm
from simtk.openmm.app import ForceField, PME, HBonds, PDBFile

import openpathsampling as paths
from openpathsampling.storage import Storage
from openpathsampling.todict import restores_as_full_object

@restores_as_full_object
class OpenMMEngine(paths.DynamicsEngine):
    """OpenMM dynamics engine."""

    units = {
        'length': u.nanometers,
        'velocity': u.nanometers / u.picoseconds,
        'energy': u.joule / u.mole
    }

    _default_options = {
        'nsteps_per_frame': 10,
        'solute_indices': [0],
        'n_frames_max': 5000,
        "temperature": 300.0 * u.kelvin,
        'collision_rate': 1.0 / u.picoseconds,
        'timestep': 2.0 * u.femtoseconds,
        'platform': 'fastest',
        'forcefield_solute': 'amber96.xml',
        'forcefield_solvent': 'tip3p.xml'
    }

    @staticmethod
    def auto(filename, template, options, mode='auto', units=None):
        """
        Create or Restore a OpenMMEngine

        Parameters
        ----------
        filename : str
            the filename of the storage
        template : Snapshot
            the template Snapshot to be used. It contains the necessary
            units as well as the topology
        options : dict of { str : str }
            a dictionary that contains the parameters used in the
            construction of the OpenMMEngine
        mode : str ('restore', 'create' or 'auto')
            a string setting the mode of creation or restoration. The option
            'auto' (default) will only create a new storage if the file does
            not exist yet
        units : dict of {str : simtk.unit.Unit } or None (default)
            representing a dict of string representing a dimension
            ('length', 'velocity', 'energy') pointing the the
            simtk.unit.Unit to be used. This overrides the units used in the
            template

        Returns
        -------
        OpenMMEngine
            the created or restored OpenMMEngine instance

        """
        if mode == 'auto':
            if os.path.isfile(filename):
                mode = 'restore'
            else:
                mode = 'create'

        if mode == 'create':
            return OpenMMEngine._create_with_storage(filename, template, options, units)
        elif mode == 'restore':
            return OpenMMEngine._restore_from_storage(filename)
        else:
            raise ValueError('Unknown mode: ' + mode)
            return None

    @staticmethod
    def _create_with_storage(filename, template, options, units=None):
        # for openmm we will create a suitable for configuration with
        # attached box_vectors and topology

        if type(template) is str:
            template = paths.snapshot_from_pdb(template, units=units)

        # once we have a template configuration (coordinates to not really
        # matter) we can create a storage. We might move this logic out of
        # the dynamics engine and keep storage and engine generation
        # completely separate!

        storage = Storage(
            filename=filename,
            template=template,
            mode='w'
        )

        # save simulator options, should be replaced by just saving the simulator object

        options['template'] = template

#        storage.init_str('simulation_options')
#        storage.write_as_json('simulation_options', options)

        engine = OpenMMEngine(
            options=options
        )
        engine.storage = storage
        storage.engine.save(engine)

        return engine

    @staticmethod
    def _restore_from_storage(filename):
        # open storage, which also gets the topology!
        storage = Storage(
            filename=filename,
            mode='a'
        )

#        options = storage.restore_object('simulation_options')

#        options['template'] = storage.template

#        engine = OpenMMEngine(
#            options=options
#        )

        engine = storage.engine.load(0)

        engine.storage = storage
        return engine

    def __init__(self, options, template=None):

        if 'template' in options:
            template = options['template']

        self.topology = template.topology
        self.options = {
        }

        super(OpenMMEngine, self).__init__(
            options=options,
            template=template
        )

        # set up the OpenMM simulation
        forcefield = ForceField( self.options["forcefield_solute"],
                                 self.options["forcefield_solvent"] )

        openmm_topology = paths.to_openmm_topology(self.template)

        system = forcefield.createSystem( openmm_topology,
                                          nonbondedMethod=PME,
                                          nonbondedCutoff=1.0 * u.nanometers,
                                          constraints=HBonds )

        integrator = VVVRIntegrator( self.options["temperature"],
                                     self.options["collision_rate"],
                                     self.options["timestep"] )

        simulation = openmm.app.Simulation(openmm_topology, system,
                                           integrator)

        # claim the OpenMM simulation as our own
        self.simulation = simulation

        # set no cached snapshot, menas it will be constructed from the openmm context
        self._current_snapshot = None
        self._current_momentum = None
        self._current_configuration = None
        self._current_box_vectors = None

    def equilibrate(self, nsteps):
        # TODO: rename... this is position restrained equil, right?
        #self.simulation.context.setPositions(self.pdb.positions) #TODO move
        system = self.simulation.system
        n_solute = len(self.solute_indices)

        solute_masses = u.Quantity(np.zeros(n_solute, np.double), u.dalton)
        for i in self.solute_indices:
            solute_masses[i] = system.getParticleMass(i)
            system.setParticleMass(i,0.0)

        self.simulation.step(nsteps)

        # empty cache
        self._current_snapshot = None

        for i in self.solute_indices:
            system.setParticleMass(i, solute_masses[i].value_in_unit(u.dalton))

    # this property is specific to direct control simulations: other
    # simulations might not use this
    # TODO: Maybe remove this and put it into the creation logic
    @property
    def nsteps_per_frame(self):
        return self._nsteps_per_frame

    @nsteps_per_frame.setter
    def nsteps_per_frame(self, value):
        self._nsteps_per_frame = value

    # TODO: there are two reasonable approaches to this:
    # 1. require that part of the engine.next_frame() function be that the
    #    user saves a snapshot object called `self._current_snapshot`
    # 2. build the current snapshot on the fly every time the snapshot is
    #    needed
    # The trade-off is that (1) will be faster if we ask for the snapshot
    # frequently, but it is also much more likely to be a source of errors
    # for users who forget to implement that last step.

    def _build_current_snapshot(self):
        # TODO: Add caching for this and mark if changed

        print 'get state'
        state = self.simulation.context.getState(getPositions=True,
                                                 getVelocities=True,
                                                 getEnergy=True)

        print 'build snapshot'
        return paths.Snapshot(coordinates = state.getPositions(asNumpy=True),
                        box_vectors = state.getPeriodicBoxVectors(asNumpy=True),
                        potential_energy = state.getPotentialEnergy(),
                        velocities = state.getVelocities(asNumpy=True),
                        kinetic_energy = state.getKineticEnergy(),
                        topology = self.topology
                       )

    @property
    def current_snapshot(self):
        if self._current_snapshot is None:
            self._current_snapshot = self._build_current_snapshot()

        return self._current_snapshot

    def _changed(self):
        self._current_snapshot = None

    @current_snapshot.setter
    def current_snapshot(self, snapshot):
        if snapshot is not self._current_snapshot:
            if snapshot.configuration is not None:
                if self._current_snapshot is None or snapshot.configuration is not self._current_snapshot.configuration:
                    # new snapshot has a different configuration so update
                    self.simulation.context.setPositions(snapshot.coordinates)

                    # TODO: Check if this is the right way to make sure the box is right!
#                    if self._current_snapshot is None or snapshot.box_vectors != self._current_snapshot.box_vectors:
#                        self.simulation.context.getPeriodicBoxVectors(snapshot.box_vectors)

            if snapshot.momentum is not None:
                if self._current_snapshot is None or snapshot.momentum is not self._current_snapshot.momentum or snapshot.reversed != self._current_snapshot.reversed:
                    # new snapshot has a different momenta (different coordinates and reverse setting)
                    # so update. Note snapshot.velocities is different from snapshot.momenta.velocities!!!
                    # The first includes the reversal setting in the snapshot the second does not.
#                    print snapshot.momentum.velocities
                    self.simulation.context.setVelocities(snapshot.velocities)

            # After the updates cache the new snapshot
            self._current_snapshot = snapshot

    def generate_next_frame(self):
        print 'run for steps ', self.nsteps_per_frame
        self.simulation.step(self.nsteps_per_frame)
        self._current_snapshot = None
        return self.current_snapshot

    # (possibly temporary) shortcuts for momentum and configuration
    @property
    def momentum(self):
        return self.current_snapshot.momentum

    # remove setters here, because it might lead to wrong velocities
    # because of incorrect treatment of reversed, it will also violate the
    # immutable character of the cached snapshot!
#    @momentum.setter
#    def momentum(self, momentum):
#        if momentum is not self._current_momentum:
#            self._current_momentum = momentum
#            self.simulation.context.setVelocities(momentum.velocities)

    @property
    def configuration(self):
        return self.current_snapshot.configuration

#    @configuration.setter
#    def configuration(self, config):
#        if config is not self._current_configuration:
#            self._current_configuration = config
#            self.simulation.context.setPositions(config.coordinates)
#
#            # TODO: Check if this is the right way to make sure the box is right!
#            if False and config.box_vectors != self._current_box_vectors:
#                self.simulation.context.getPeriodicBoxVectors(config.box_vectors)
#                self._current_box_vectors = config.box_vectors

<commit_msg>close file after pdb read<commit_after>import os
import numpy as np
from openmmtools.integrators import VVVRIntegrator

import simtk.unit as u
import simtk.openmm as openmm
from simtk.openmm.app import ForceField, PME, HBonds, PDBFile

import openpathsampling as paths
from openpathsampling.storage import Storage
from openpathsampling.todict import restores_as_full_object

@restores_as_full_object
class OpenMMEngine(paths.DynamicsEngine):
    """OpenMM dynamics engine."""

    units = {
        'length': u.nanometers,
        'velocity': u.nanometers / u.picoseconds,
        'energy': u.joule / u.mole
    }

    _default_options = {
        'nsteps_per_frame': 10,
        'solute_indices': [0],
        'n_frames_max': 5000,
        "temperature": 300.0 * u.kelvin,
        'collision_rate': 1.0 / u.picoseconds,
        'timestep': 2.0 * u.femtoseconds,
        'platform': 'fastest',
        'forcefield_solute': 'amber96.xml',
        'forcefield_solvent': 'tip3p.xml'
    }

    @staticmethod
    def auto(filename, template, options, mode='auto', units=None):
        """
        Create or Restore a OpenMMEngine

        Parameters
        ----------
        filename : str
            the filename of the storage
        template : Snapshot
            the template Snapshot to be used. It contains the necessary
            units as well as the topology
        options : dict of { str : str }
            a dictionary that contains the parameters used in the
            construction of the OpenMMEngine
        mode : str ('restore', 'create' or 'auto')
            a string setting the mode of creation or restoration. The option
            'auto' (default) will only create a new storage if the file does
            not exist yet
        units : dict of {str : simtk.unit.Unit } or None (default)
            representing a dict of string representing a dimension
            ('length', 'velocity', 'energy') pointing the the
            simtk.unit.Unit to be used. This overrides the units used in the
            template

        Returns
        -------
        OpenMMEngine
            the created or restored OpenMMEngine instance

        """
        if mode == 'auto':
            if os.path.isfile(filename):
                mode = 'restore'
            else:
                mode = 'create'

        if mode == 'create':
            return OpenMMEngine._create_with_storage(filename, template, options, units)
        elif mode == 'restore':
            return OpenMMEngine._restore_from_storage(filename)
        else:
            raise ValueError('Unknown mode: ' + mode)
            return None

    @staticmethod
    def _create_with_storage(filename, template, options, units=None):
        # for openmm we will create a suitable for configuration with
        # attached box_vectors and topology

        if type(template) is str:
            template = paths.snapshot_from_pdb(template, units=units)

        # once we have a template configuration (coordinates to not really
        # matter) we can create a storage. We might move this logic out of
        # the dynamics engine and keep storage and engine generation
        # completely separate!

        storage = Storage(
            filename=filename,
            template=template,
            mode='w'
        )

        # save simulator options, should be replaced by just saving the simulator object

        options['template'] = template

#        storage.init_str('simulation_options')
#        storage.write_as_json('simulation_options', options)

        engine = OpenMMEngine(
            options=options
        )
        engine.storage = storage
        storage.engine.save(engine)

        return engine

    @staticmethod
    def _restore_from_storage(filename):
        # open storage, which also gets the topology!
        storage = Storage(
            filename=filename,
            mode='a'
        )

#        options = storage.restore_object('simulation_options')

#        options['template'] = storage.template

#        engine = OpenMMEngine(
#            options=options
#        )

        engine = storage.engine.load(0)

        engine.storage = storage
        return engine

    def __init__(self, options, template=None):

        if 'template' in options:
            template = options['template']

        self.topology = template.topology
        self.options = {
        }

        super(OpenMMEngine, self).__init__(
            options=options,
            template=template
        )

        # set up the OpenMM simulation
        forcefield = ForceField( self.options["forcefield_solute"],
                                 self.options["forcefield_solvent"] )

        openmm_topology = paths.to_openmm_topology(self.template)

        system = forcefield.createSystem( openmm_topology,
                                          nonbondedMethod=PME,
                                          nonbondedCutoff=1.0 * u.nanometers,
                                          constraints=HBonds )

        integrator = VVVRIntegrator( self.options["temperature"],
                                     self.options["collision_rate"],
                                     self.options["timestep"] )

        simulation = openmm.app.Simulation(openmm_topology, system,
                                           integrator)

        # claim the OpenMM simulation as our own
        self.simulation = simulation

        # set no cached snapshot, menas it will be constructed from the openmm context
        self._current_snapshot = None
        self._current_momentum = None
        self._current_configuration = None
        self._current_box_vectors = None

    def equilibrate(self, nsteps):
        # TODO: rename... this is position restrained equil, right?
        #self.simulation.context.setPositions(self.pdb.positions) #TODO move
        system = self.simulation.system
        n_solute = len(self.solute_indices)

        solute_masses = u.Quantity(np.zeros(n_solute, np.double), u.dalton)
        for i in self.solute_indices:
            solute_masses[i] = system.getParticleMass(i)
            system.setParticleMass(i,0.0)

        self.simulation.step(nsteps)

        # empty cache
        self._current_snapshot = None

        for i in self.solute_indices:
            system.setParticleMass(i, solute_masses[i].value_in_unit(u.dalton))

    # this property is specific to direct control simulations: other
    # simulations might not use this
    # TODO: Maybe remove this and put it into the creation logic
    @property
    def nsteps_per_frame(self):
        return self._nsteps_per_frame

    @nsteps_per_frame.setter
    def nsteps_per_frame(self, value):
        self._nsteps_per_frame = value

    # TODO: there are two reasonable approaches to this:
    # 1. require that part of the engine.next_frame() function be that the
    #    user saves a snapshot object called `self._current_snapshot`
    # 2. build the current snapshot on the fly every time the snapshot is
    #    needed
    # The trade-off is that (1) will be faster if we ask for the snapshot
    # frequently, but it is also much more likely to be a source of errors
    # for users who forget to implement that last step.

    def _build_current_snapshot(self):
        # TODO: Add caching for this and mark if changed

        print 'get state'
        state = self.simulation.context.getState(getPositions=True,
                                                 getVelocities=True,
                                                 getEnergy=True)

        print 'build snapshot'
        return paths.Snapshot(coordinates = state.getPositions(asNumpy=True),
                        box_vectors = state.getPeriodicBoxVectors(asNumpy=True),
                        potential_energy = state.getPotentialEnergy(),
                        velocities = state.getVelocities(asNumpy=True),
                        kinetic_energy = state.getKineticEnergy(),
                        topology = self.topology
                       )

    @property
    def current_snapshot(self):
        if self._current_snapshot is None:
            self._current_snapshot = self._build_current_snapshot()

        return self._current_snapshot

    def _changed(self):
        self._current_snapshot = None

    @current_snapshot.setter
    def current_snapshot(self, snapshot):
        if snapshot is not self._current_snapshot:
            if snapshot.configuration is not None:
                if self._current_snapshot is None or snapshot.configuration is not self._current_snapshot.configuration:
                    # new snapshot has a different configuration so update
                    self.simulation.context.setPositions(snapshot.coordinates)

                    # TODO: Check if this is the right way to make sure the box is right!
#                    if self._current_snapshot is None or snapshot.box_vectors != self._current_snapshot.box_vectors:
#                        self.simulation.context.getPeriodicBoxVectors(snapshot.box_vectors)

            if snapshot.momentum is not None:
                if self._current_snapshot is None or snapshot.momentum is not self._current_snapshot.momentum or snapshot.reversed != self._current_snapshot.reversed:
                    # new snapshot has a different momenta (different coordinates and reverse setting)
                    # so update. Note snapshot.velocities is different from snapshot.momenta.velocities!!!
                    # The first includes the reversal setting in the snapshot the second does not.
#                    print snapshot.momentum.velocities
                    self.simulation.context.setVelocities(snapshot.velocities)

            # After the updates cache the new snapshot
            self._current_snapshot = snapshot

    def generate_next_frame(self):
        print 'run for steps ', self.nsteps_per_frame
        self.simulation.step(self.nsteps_per_frame)
        self._current_snapshot = None
        print 'finished steps'
        return self.current_snapshot

    # (possibly temporary) shortcuts for momentum and configuration
    @property
    def momentum(self):
        return self.current_snapshot.momentum

    # remove setters here, because it might lead to wrong velocities
    # because of incorrect treatment of reversed, it will also violate the
    # immutable character of the cached snapshot!
#    @momentum.setter
#    def momentum(self, momentum):
#        if momentum is not self._current_momentum:
#            self._current_momentum = momentum
#            self.simulation.context.setVelocities(momentum.velocities)

    @property
    def configuration(self):
        return self.current_snapshot.configuration

#    @configuration.setter
#    def configuration(self, config):
#        if config is not self._current_configuration:
#            self._current_configuration = config
#            self.simulation.context.setPositions(config.coordinates)
#
#            # TODO: Check if this is the right way to make sure the box is right!
#            if False and config.box_vectors != self._current_box_vectors:
#                self.simulation.context.getPeriodicBoxVectors(config.box_vectors)
#                self._current_box_vectors = config.box_vectors

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2018 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import collections
import logging
import numpy

from openquake.hazardlib.calc.gmf import GmfComputer
from openquake.hazardlib.gsim.base import ContextMaker
from openquake.commonlib import readinput, source, calc
from openquake.hazardlib.source.rupture import EBRupture
from openquake.calculators import base

@base.calculators.add('scenario')
class ScenarioCalculator(base.HazardCalculator):
    """
    Scenario hazard calculator
    """
    is_stochastic = True

    def pre_execute(self):
        """
        Read the site collection and initialize GmfComputer and seeds
        """
        oq = self.oqparam
        cinfo = source.CompositionInfo.fake(readinput.get_gsim_lt(oq))
        self.datastore['csm_info'] = cinfo
        if 'rupture_model' not in oq.inputs:
            logging.warn('There is no rupture_model, the calculator will just '
                         'import data without performing any calculation')
            super().pre_execute()
            return
        self.rup = readinput.get_rupture(oq)
        self.gsims = readinput.get_gsims(oq)
        self.cmaker = ContextMaker(self.gsims, oq.maximum_distance,
                                   {'filter_distance': oq.filter_distance})
        super().pre_execute()
        self.datastore['oqparam'] = oq
        self.rlzs_assoc = cinfo.get_rlzs_assoc()
        E = oq.number_of_ground_motion_fields
        events = numpy.zeros(E, readinput.stored_event_dt)
        events['eid'] = numpy.arange(E)
        ebr = EBRupture(self.rup, 0, 0, self.sitecol.sids, numpy.array([E]))
        self.datastore['events'] = events
        rupser = calc.RuptureSerializer(self.datastore)
        rupser.save([ebr])
        rupser.close()
        self.computer = GmfComputer(
            ebr, self.sitecol, oq.imtls, self.cmaker, oq.truncation_level,
            oq.correl_model)

    def init(self):
        pass

    def execute(self):
        """
        Compute the GMFs and return a dictionary gsim -> array(N, E, I)
        """
        self.gmfa = collections.OrderedDict()
        if 'rupture_model' not in self.oqparam.inputs:
            return self.gmfa
        with self.monitor('computing gmfs'):
            E = self.oqparam.number_of_ground_motion_fields
            for gsim in self.gsims:
                gmfa = self.computer.compute(gsim, E)  # shape (I, N, E)
                self.gmfa[gsim] = gmfa.transpose(1, 2, 0)  # shape (N, E, I)
        return self.gmfa

    def post_execute(self, dummy):
        if self.gmfa:
            with self.monitor('saving gmfs', autoflush=True):
                base.save_gmf_data(
                    self.datastore, self.sitecol,
                    numpy.array(list(self.gmfa.values())), self.oqparam.imtls)
<commit_msg>[skip hazardlib]<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2018 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import collections
import logging
import numpy

from openquake.hazardlib.calc.gmf import GmfComputer
from openquake.hazardlib.gsim.base import ContextMaker
from openquake.commonlib import readinput, source, calc
from openquake.hazardlib.source.rupture import EBRupture
from openquake.calculators import base

@base.calculators.add('scenario')
class ScenarioCalculator(base.HazardCalculator):
    """
    Scenario hazard calculator
    """
    is_stochastic = True

    def pre_execute(self):
        """
        Read the site collection and initialize GmfComputer and seeds
        """
        oq = self.oqparam
        cinfo = source.CompositionInfo.fake(readinput.get_gsim_lt(oq))
        self.datastore['csm_info'] = cinfo
        if 'rupture_model' not in oq.inputs:
            logging.warn('There is no rupture_model, the calculator will just '
                         'import data without performing any calculation')
            super().pre_execute()
            return
        self.rup = readinput.get_rupture(oq)
        self.gsims = readinput.get_gsims(oq)
        self.cmaker = ContextMaker(self.gsims, oq.maximum_distance,
                                   {'filter_distance': oq.filter_distance})
        super().pre_execute()
        self.datastore['oqparam'] = oq
        self.rlzs_assoc = cinfo.get_rlzs_assoc()
        E = oq.number_of_ground_motion_fields
        events = numpy.zeros(E, readinput.stored_event_dt)
        events['eid'] = numpy.arange(E)
        ebr = EBRupture(self.rup, 0, 0, self.sitecol.sids, numpy.array([E]))
        self.datastore['events'] = events
        rupser = calc.RuptureSerializer(self.datastore)
        rupser.save([ebr], len(self.gsims))
        rupser.close()
        self.computer = GmfComputer(
            ebr, self.sitecol, oq.imtls, self.cmaker, oq.truncation_level,
            oq.correl_model)

    def init(self):
        pass

    def execute(self):
        """
        Compute the GMFs and return a dictionary gsim -> array(N, E, I)
        """
        self.gmfa = collections.OrderedDict()
        if 'rupture_model' not in self.oqparam.inputs:
            return self.gmfa
        with self.monitor('computing gmfs'):
            E = self.oqparam.number_of_ground_motion_fields
            for gsim in self.gsims:
                gmfa = self.computer.compute(gsim, E)  # shape (I, N, E)
                self.gmfa[gsim] = gmfa.transpose(1, 2, 0)  # shape (N, E, I)
        return self.gmfa

    def post_execute(self, dummy):
        if self.gmfa:
            with self.monitor('saving gmfs', autoflush=True):
                base.save_gmf_data(
                    self.datastore, self.sitecol,
                    numpy.array(list(self.gmfa.values())), self.oqparam.imtls)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2020 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import logging
import numpy

from openquake.hazardlib.calc.gmf import GmfComputer
from openquake.hazardlib.gsim.base import ContextMaker
from openquake.hazardlib.calc.stochastic import get_rup_array
from openquake.hazardlib.source.rupture import EBRupture, events_dt
from openquake.commonlib import readinput, logictree, calc
from openquake.calculators import base

@base.calculators.add('scenario')
class ScenarioCalculator(base.HazardCalculator):
    """
    Scenario hazard calculator
    """
    is_stochastic = True

    def pre_execute(self):
        """
        Read the site collection and initialize GmfComputer and seeds
        """
        oq = self.oqparam
        cinfo = logictree.FullLogicTree.fake(readinput.get_gsim_lt(oq))
        self.realizations = cinfo.get_realizations()
        self.datastore['full_lt'] = cinfo
        if 'rupture_model' not in oq.inputs:
            logging.warning(
                'There is no rupture_model, the calculator will just '
                'import data without performing any calculation')
            super().pre_execute()
            return
        self.rup = readinput.get_rupture(oq)
        self.gsims = readinput.get_gsims(oq)
        R = len(self.gsims)
        self.cmaker = ContextMaker('*', self.gsims,
                                   {'maximum_distance': oq.maximum_distance,
                                    'filter_distance': oq.filter_distance})
        super().pre_execute()
        self.datastore['oqparam'] = oq
        self.store_rlz_info({})
        rlzs_by_gsim = cinfo.get_rlzs_by_gsim(0)
        E = oq.number_of_ground_motion_fields
        n_occ = numpy.array([E])
        ebr = EBRupture(self.rup, 0, 0, n_occ)
        ebr.e0 = 0
        events = numpy.zeros(E * R, events_dt)
        for rlz, eids in ebr.get_eids_by_rlz(rlzs_by_gsim).items():
            events[rlz * E: rlz * E + E]['id'] = eids
            events[rlz * E: rlz * E + E]['rlz_id'] = rlz
        self.datastore['events'] = self.events = events
        rupser = calc.RuptureSerializer(self.datastore)
        rup_array = get_rup_array([ebr], self.src_filter())
        if len(rup_array) == 0:
            maxdist = oq.maximum_distance(
                self.rup.tectonic_region_type, self.rup.mag)
            raise RuntimeError('There are no sites within the maximum_distance'
                               ' of %s km from the rupture' % maxdist)
        rupser.save(rup_array)
        rupser.close()
        self.computer = GmfComputer(
            ebr, self.sitecol, oq.imtls, self.cmaker, oq.truncation_level,
            oq.correl_model, self.amplifier)
        M32 = (numpy.float32, (len(self.oqparam.imtls),))
        self.sig_eps_dt = [('eid', numpy.uint64), ('sig', M32), ('eps', M32)]

    def init(self):
        pass

    def execute(self):
        """
        Compute the GMFs and return a dictionary gsim -> array(N, E, I)
        """
        arrays = []
        if 'rupture_model' not in self.oqparam.inputs:
            return ()
        n = self.oqparam.number_of_ground_motion_fields
        with self.monitor('computing gmfs'):
            ei = 0
            for gsim in self.gsims:
                gmfa, sig, eps = self.computer.compute(gsim, n)
                lst = []
                for s, e in zip(sig.T, eps.T):  # shape (M, E) -> (E, M)
                    lst.append((ei, s, e))
                    ei += 1
                arrays.append(gmfa.transpose(1, 2, 0))  # shape (N, n, I)
        self.datastore['gmf_data/sigma_epsilon'] = numpy.array(
            lst, self.sig_eps_dt)
        return numpy.concatenate(arrays, axis=1)  # shape (N, E, I)

    def post_execute(self, gmfa):
        if len(gmfa) == 0:  # no rupture_model
            return
        with self.monitor('saving gmfs'):
            base.save_gmf_data(
                self.datastore, self.sitecol, gmfa,
                self.oqparam.imtls, self.events)
<commit_msg>Small renaming [skip CI]<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2020 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import logging
import numpy

from openquake.hazardlib.calc.gmf import GmfComputer
from openquake.hazardlib.gsim.base import ContextMaker
from openquake.hazardlib.calc.stochastic import get_rup_array
from openquake.hazardlib.source.rupture import EBRupture, events_dt
from openquake.commonlib import readinput, logictree, calc
from openquake.calculators import base

@base.calculators.add('scenario')
class ScenarioCalculator(base.HazardCalculator):
    """
    Scenario hazard calculator
    """
    is_stochastic = True

    def pre_execute(self):
        """
        Read the site collection and initialize GmfComputer and seeds
        """
        oq = self.oqparam
        full_lt = logictree.FullLogicTree.fake(readinput.get_gsim_lt(oq))
        self.realizations = full_lt.get_realizations()
        self.datastore['full_lt'] = full_lt
        if 'rupture_model' not in oq.inputs:
            logging.warning(
                'There is no rupture_model, the calculator will just '
                'import data without performing any calculation')
            super().pre_execute()
            return
        self.rup = readinput.get_rupture(oq)
        self.gsims = readinput.get_gsims(oq)
        R = len(self.gsims)
        self.cmaker = ContextMaker('*', self.gsims,
                                   {'maximum_distance': oq.maximum_distance,
                                    'filter_distance': oq.filter_distance})
        super().pre_execute()
        self.datastore['oqparam'] = oq
        self.store_rlz_info({})
        rlzs_by_gsim = full_lt.get_rlzs_by_gsim(0)
        E = oq.number_of_ground_motion_fields
        n_occ = numpy.array([E])
        ebr = EBRupture(self.rup, 0, 0, n_occ)
        ebr.e0 = 0
        events = numpy.zeros(E * R, events_dt)
        for rlz, eids in ebr.get_eids_by_rlz(rlzs_by_gsim).items():
            events[rlz * E: rlz * E + E]['id'] = eids
            events[rlz * E: rlz * E + E]['rlz_id'] = rlz
        self.datastore['events'] = self.events = events
        rupser = calc.RuptureSerializer(self.datastore)
        rup_array = get_rup_array([ebr], self.src_filter())
        if len(rup_array) == 0:
            maxdist = oq.maximum_distance(
                self.rup.tectonic_region_type, self.rup.mag)
            raise RuntimeError('There are no sites within the maximum_distance'
                               ' of %s km from the rupture' % maxdist)
        rupser.save(rup_array)
        rupser.close()
        self.computer = GmfComputer(
            ebr, self.sitecol, oq.imtls, self.cmaker, oq.truncation_level,
            oq.correl_model, self.amplifier)
        M32 = (numpy.float32, (len(self.oqparam.imtls),))
        self.sig_eps_dt = [('eid', numpy.uint64), ('sig', M32), ('eps', M32)]

    def init(self):
        pass

    def execute(self):
        """
        Compute the GMFs and return a dictionary gsim -> array(N, E, I)
        """
        arrays = []
        if 'rupture_model' not in self.oqparam.inputs:
            return ()
        n = self.oqparam.number_of_ground_motion_fields
        with self.monitor('computing gmfs'):
            ei = 0
            for gsim in self.gsims:
                gmfa, sig, eps = self.computer.compute(gsim, n)
                lst = []
                for s, e in zip(sig.T, eps.T):  # shape (M, E) -> (E, M)
                    lst.append((ei, s, e))
                    ei += 1
                arrays.append(gmfa.transpose(1, 2, 0))  # shape (N, n, I)
        self.datastore['gmf_data/sigma_epsilon'] = numpy.array(
            lst, self.sig_eps_dt)
        return numpy.concatenate(arrays, axis=1)  # shape (N, E, I)

    def post_execute(self, gmfa):
        if len(gmfa) == 0:  # no rupture_model
            return
        with self.monitor('saving gmfs'):
            base.save_gmf_data(
                self.datastore, self.sitecol, gmfa,
                self.oqparam.imtls, self.events)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2010-2018 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import re
import logging
import operator
import collections
import numpy
from openquake.hazardlib import probability_map

MAX_INT = 2 ** 31 - 1
U16 = numpy.uint16
U32 = numpy.uint32
F32 = numpy.float32
weight = operator.attrgetter('weight')
rlz_dt = numpy.dtype([
    ('branch_path', 'S200'), ('gsims', 'S100'), ('weight', F32)])

class LtRealization(object):
    """
    Composite realization build on top of a source model realization and
    a GSIM realization.
    """
    def __init__(self, ordinal, sm_lt_path, gsim_rlz, weight):
        self.ordinal = ordinal
        self.sm_lt_path = tuple(sm_lt_path)
        self.gsim_rlz = gsim_rlz
        self.weight = weight

    def __repr__(self):
        return '<%d,%s,w=%s>' % (self.ordinal, self.uid, self.weight)

    @property
    def gsim_lt_path(self):
        return self.gsim_rlz.lt_path

    @property
    def uid(self):
        """An unique identifier for effective realizations"""
        return '_'.join(self.sm_lt_path) + '~' + self.gsim_rlz.uid

    def __lt__(self, other):
        return self.ordinal < other.ordinal

    def __eq__(self, other):
        return repr(self) == repr(other)

    def __ne__(self, other):
        return repr(self) != repr(other)

    def __hash__(self):
        return hash(repr(self))

class RlzsAssoc(object):
    """
    Realization association class. It should not be instantiated directly,
    but only via the method :meth:
    `openquake.commonlib.source.CompositeSourceModel.get_rlzs_assoc`.

    :attr realizations: list of :class:`LtRealization` objects
    :attr gsim_by_trt: list of dictionaries {trt: gsim}
    :attr rlzs_assoc: dictionary {src_group_id, gsim: rlzs}
    :attr rlzs_by_smodel: list of lists of realizations

    For instance, for the non-trivial logic tree in
    :mod:`openquake.qa_tests_data.classical.case_15`, which has 4 tectonic
    region types and 4 + 2 + 2 realizations, there are the following
    associations:

    (0, 'BooreAtkinson2008()') ['#0-SM1-BA2008_C2003', '#1-SM1-BA2008_T2002']
    (0, 'CampbellBozorgnia2008()') ['#2-SM1-CB2008_C2003', '#3-SM1-CB2008_T2002']
    (1, 'Campbell2003()') ['#0-SM1-BA2008_C2003', '#2-SM1-CB2008_C2003']
    (1, 'ToroEtAl2002()') ['#1-SM1-BA2008_T2002', '#3-SM1-CB2008_T2002']
    (2, 'BooreAtkinson2008()') ['#4-SM2_a3pt2b0pt8-BA2008']
    (2, 'CampbellBozorgnia2008()') ['#5-SM2_a3pt2b0pt8-CB2008']
    (3, 'BooreAtkinson2008()') ['#6-SM2_a3b1-BA2008']
    (3, 'CampbellBozorgnia2008()') ['#7-SM2_a3b1-CB2008']
    """
    def __init__(self, csm_info):
        self.seed = csm_info.seed
        self.csm_info = csm_info
        self.num_samples = csm_info.num_samples
        self.gsim_by_trt = []  # rlz.ordinal -> {trt: gsim}
        self.rlzs_by_smodel = {sm.ordinal: [] for sm in csm_info.source_models}

    def get_rlzs_by_gsim(self, trt_or_grp_id, sm_id=None):
        """
        :param trt_or_grp_id: a tectonic region type or a source group ID
        :param sm_id: source model ordinal (or None)
        :returns: a dictionary gsim -> rlzs
        """
        if isinstance(trt_or_grp_id, (int, U16, U32)):  # grp_id
            trt = self.csm_info.trt_by_grp[trt_or_grp_id]
            sm_id = self.csm_info.get_sm_by_grp()[trt_or_grp_id]
        else:  # assume TRT string
            trt = trt_or_grp_id
        acc = collections.defaultdict(list)
        if sm_id is None:  # full dictionary
            for rlz, gsim_by_trt in zip(self.realizations, self.gsim_by_trt):
                acc[gsim_by_trt[trt]].append(rlz.ordinal)
        else:  # dictionary for the selected source model
            for rlz in self.rlzs_by_smodel[sm_id]:
                gsim_by_trt = self.gsim_by_trt[rlz.ordinal]
                try:  # if there is a single TRT
                    [gsim] = gsim_by_trt.values()
                except ValueError:  # there is more than 1 TRT
                    gsim = gsim_by_trt[trt]
                acc[gsim].append(rlz.ordinal)
        return {gsim: numpy.array(acc[gsim], dtype=U16)
                for gsim in sorted(acc)}

    def by_grp(self):
        """
        :returns: a dictionary grp -> rlzis
        """
        dic = {}  # grp -> [(gsim_idx, rlzis), ...]
        for sm in self.csm_info.source_models:
            for sg in sm.src_groups:
                if not sg.eff_ruptures:
                    continue
                rlzs_by_gsim = self.get_rlzs_by_gsim(sg.trt, sm.ordinal)
                if not rlzs_by_gsim:
                    continue
                dic['grp-%02d' % sg.id] = numpy.array(
                    list(rlzs_by_gsim.values()))
        return dic

    def _init(self):
        """
        Finalize the initialization of the RlzsAssoc object by setting
        the (reduced) weights of the realizations.
        """
        if self.num_samples:
            assert len(self.realizations) == self.num_samples, (
                len(self.realizations), self.num_samples)
            tot_weight = sum(rlz.weight for rlz in self.realizations)
            for rlz in self.realizations:
                rlz.weight /= tot_weight
        else:
            tot_weight = sum(rlz.weight for rlz in self.realizations)
            if not tot_weight.is_one():
                # this may happen for rounding errors or because of the
                # logic tree reduction; we ensure the sum of the weights is 1
                for rlz in self.realizations:
                    rlz.weight = rlz.weight / tot_weight

    @property
    def realizations(self):
        """Flat list with all the realizations"""
        return sum(self.rlzs_by_smodel.values(), [])

    @property
    def weights(self):
        """Array with the weight of the realizations"""
        return numpy.array([rlz.weight for rlz in self.realizations])

    def combine_pmaps(self, pmap_by_grp):
        """
        :param pmap_by_grp: dictionary group string -> probability map
        :returns: a list of probability maps, one per realization
        """
        grp = list(pmap_by_grp)[0]  # pmap_by_grp must be non-empty
        num_levels = pmap_by_grp[grp].shape_y
        pmaps = [probability_map.ProbabilityMap(num_levels, 1)
                 for _ in self.realizations]
        array = self.by_grp()
        for grp in pmap_by_grp:
            for gsim_idx, rlzis in enumerate(array[grp]):
                pmap = pmap_by_grp[grp].extract(gsim_idx)
                for rlzi in rlzis:
                    pmaps[rlzi] |= pmap
        return pmaps

    def get_rlz(self, rlzstr):
        """
        Get a Realization instance for a string of the form 'rlz-\d+'
        """
        mo = re.match('rlz-(\d+)', rlzstr)
        if not mo:
            return
        return self.realizations[int(mo.group(1))]

    def _add_realizations(self, offset, lt_model, all_trts, gsim_rlzs):
        idx = numpy.arange(offset, offset + len(gsim_rlzs))
        rlzs = []
        for i, gsim_rlz in enumerate(gsim_rlzs):
            weight = float(lt_model.weight) * gsim_rlz.weight
            rlz = LtRealization(idx[i], lt_model.path, gsim_rlz, weight)
            self.gsim_by_trt.append(dict(zip(all_trts, gsim_rlz.value)))
            rlzs.append(rlz)
        self.rlzs_by_smodel[lt_model.ordinal] = rlzs

    def __repr__(self):
        pairs = []
        dic = self.by_grp()
        size = sum(len(dic[grp]) for grp in dic)
        for grp in sorted(dic):
            grp_id = int(grp[4:])
            gsims = self.csm_info.get_gsims(grp_id)
            for gsim_idx, rlzis in enumerate(dic[grp]):
                if len(rlzis) > 10:  # short representation
                    rlzis = ['%d realizations' % len(rlzis)]
                pairs.append(('%s,%s' % (grp_id, gsims[gsim_idx]), rlzis))
        return '<%s(size=%d, rlzs=%d)\n%s>' % (
            self.__class__.__name__, size, len(self.realizations),
            '\n'.join('%s: %s' % pair for pair in pairs))

def accept_path(path, ref_path):
    """
    :param path: a logic tree path (list or tuple of strings)
    :param ref_path: reference logic tree path
    :returns: True if `path` is consistent with `ref_path`, False otherwise

    >>> accept_path(['SM2'], ('SM2', 'a3b1'))
    False
    >>> accept_path(['SM2', '@'], ('SM2', 'a3b1'))
    True
    >>> accept_path(['@', 'a3b1'], ('SM2', 'a3b1'))
    True
    >>> accept_path('@@', ('SM2', 'a3b1'))
    True
    """
    if len(path) != len(ref_path):
        return False
    for a, b in zip(path, ref_path):
        if a != '@' and a != b:
            return False
    return True

def get_rlzs_assoc(cinfo, sm_lt_path=None, trts=None):
    """
    :param cinfo: a :class:`openquake.commonlib.source.CompositionInfo`
    :param sm_lt_path: logic tree path tuple used to select a source model
    :param trts: tectonic region types to accept
    """
    assoc = RlzsAssoc(cinfo)
    offset = 0
    trtset = set(cinfo.gsim_lt.tectonic_region_types)
    for smodel in cinfo.source_models:
        # discard source models with non-acceptable lt_path
        if sm_lt_path and not accept_path(smodel.path, sm_lt_path):
            continue

        # collect the effective tectonic region types and ruptures
        trts_ = set()
        for sg in smodel.src_groups:
            if sg.eff_ruptures:
                if (trts and sg.trt in trts) or not trts:
                    trts_.add(sg.trt)

        # recompute the GSIM logic tree if needed
        if trtset != trts_:
            before = cinfo.gsim_lt.get_num_paths()
            gsim_lt = cinfo.gsim_lt.reduce(trts_)
            after = gsim_lt.get_num_paths()
            if sm_lt_path and before > after:
                # print the warning only when saving the logic tree,
                # i.e. when called with sm_lt_path in store_rlz_info
                logging.warning('Reducing the logic tree of %s from %d to %d '
                             'realizations', smodel.name, before, after)
            gsim_rlzs = list(gsim_lt)
            all_trts = gsim_lt.all_trts
        else:
            gsim_rlzs = cinfo.gsim_rlzs
            all_trts = cinfo.gsim_lt.all_trts

        rlzs = cinfo._get_rlzs(smodel, gsim_rlzs, cinfo.seed + offset)
        assoc._add_realizations(offset, smodel, all_trts, rlzs)
        offset += len(rlzs)

    if assoc.realizations:
        assoc._init()
    return assoc
<commit_msg>[skip hazardlib]<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2010-2018 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import re
import logging
import operator
import collections
import numpy
from openquake.hazardlib import probability_map

MAX_INT = 2 ** 31 - 1
U16 = numpy.uint16
U32 = numpy.uint32
F32 = numpy.float32
weight = operator.attrgetter('weight')
rlz_dt = numpy.dtype([
    ('branch_path', 'S200'), ('gsims', 'S100'), ('weight', F32)])

class LtRealization(object):
    """
    Composite realization build on top of a source model realization and
    a GSIM realization.
    """
    def __init__(self, ordinal, sm_lt_path, gsim_rlz, weight):
        self.ordinal = ordinal
        self.sm_lt_path = tuple(sm_lt_path)
        self.gsim_rlz = gsim_rlz
        self.weight = weight

    def __repr__(self):
        return '<%d,%s,w=%s>' % (self.ordinal, self.uid, self.weight)

    @property
    def gsim_lt_path(self):
        return self.gsim_rlz.lt_path

    @property
    def uid(self):
        """An unique identifier for effective realizations"""
        return '_'.join(self.sm_lt_path) + '~' + self.gsim_rlz.uid

    def __lt__(self, other):
        return self.ordinal < other.ordinal

    def __eq__(self, other):
        return repr(self) == repr(other)

    def __ne__(self, other):
        return repr(self) != repr(other)

    def __hash__(self):
        return hash(repr(self))

class RlzsAssoc(object):
    """
    Realization association class. It should not be instantiated directly,
    but only via the method :meth:
    `openquake.commonlib.source.CompositeSourceModel.get_rlzs_assoc`.

    :attr realizations: list of :class:`LtRealization` objects
    :attr gsim_by_trt: list of dictionaries {trt: gsim}
    :attr rlzs_assoc: dictionary {src_group_id, gsim: rlzs}
    :attr rlzs_by_smodel: list of lists of realizations

    For instance, for the non-trivial logic tree in
    :mod:`openquake.qa_tests_data.classical.case_15`, which has 4 tectonic
    region types and 4 + 2 + 2 realizations, there are the following
    associations:

    (0, 'BooreAtkinson2008()') ['#0-SM1-BA2008_C2003', '#1-SM1-BA2008_T2002']
    (0, 'CampbellBozorgnia2008()') ['#2-SM1-CB2008_C2003', '#3-SM1-CB2008_T2002']
    (1, 'Campbell2003()') ['#0-SM1-BA2008_C2003', '#2-SM1-CB2008_C2003']
    (1, 'ToroEtAl2002()') ['#1-SM1-BA2008_T2002', '#3-SM1-CB2008_T2002']
    (2, 'BooreAtkinson2008()') ['#4-SM2_a3pt2b0pt8-BA2008']
    (2, 'CampbellBozorgnia2008()') ['#5-SM2_a3pt2b0pt8-CB2008']
    (3, 'BooreAtkinson2008()') ['#6-SM2_a3b1-BA2008']
    (3, 'CampbellBozorgnia2008()') ['#7-SM2_a3b1-CB2008']
    """
    def __init__(self, csm_info):
        self.seed = csm_info.seed
        self.csm_info = csm_info
        self.num_samples = csm_info.num_samples
        self.gsim_by_trt = []  # rlz.ordinal -> {trt: gsim}
        self.rlzs_by_smodel = {sm.ordinal: [] for sm in csm_info.source_models}

    def get_rlzs_by_gsim(self, trt_or_grp_id, sm_id=None):
        """
        :param trt_or_grp_id: a tectonic region type or a source group ID
        :param sm_id: source model ordinal (or None)
        :returns: a dictionary gsim -> rlzs
        """
        if isinstance(trt_or_grp_id, (int, U16, U32)):  # grp_id
            trt = self.csm_info.trt_by_grp[trt_or_grp_id]
            sm_id = self.csm_info.get_sm_by_grp()[trt_or_grp_id]
        else:  # assume TRT string
            trt = trt_or_grp_id
        acc = collections.defaultdict(list)
        if sm_id is None:  # full dictionary
            for rlz, gsim_by_trt in zip(self.realizations, self.gsim_by_trt):
                acc[gsim_by_trt[trt]].append(rlz.ordinal)
        else:  # dictionary for the selected source model
            for rlz in self.rlzs_by_smodel[sm_id]:
                gsim_by_trt = self.gsim_by_trt[rlz.ordinal]
                try:  # if there is a single TRT
                    [gsim] = gsim_by_trt.values()
                except ValueError:  # there is more than 1 TRT
                    gsim = gsim_by_trt[trt]
                acc[gsim].append(rlz.ordinal)
        return {gsim: numpy.array(acc[gsim], dtype=U16)
                for gsim in sorted(acc)}

    def by_grp(self):
        """
        :returns: a dictionary grp -> rlzis
        """
        dic = {}  # grp -> [(gsim_idx, rlzis), ...]
        for sm in self.csm_info.source_models:
            for sg in sm.src_groups:
                if not sg.eff_ruptures:
                    continue
                rlzs_by_gsim = self.get_rlzs_by_gsim(sg.trt, sm.ordinal)
                if not rlzs_by_gsim:
                    continue
                dic['grp-%02d' % sg.id] = numpy.array(
                    list(rlzs_by_gsim.values()))
        return dic

    def _init(self):
        """
        Finalize the initialization of the RlzsAssoc object by setting
        the (reduced) weights of the realizations.
        """
        if self.num_samples:
            assert len(self.realizations) == self.num_samples, (
                len(self.realizations), self.num_samples)
            tot_weight = sum(rlz.weight for rlz in self.realizations)
            for rlz in self.realizations:
                rlz.weight /= tot_weight
        else:
            tot_weight = sum(rlz.weight for rlz in self.realizations)
            if not tot_weight.is_one():
                # this may happen for rounding errors or because of the
                # logic tree reduction; we ensure the sum of the weights is 1
                for rlz in self.realizations:
                    rlz.weight = rlz.weight / tot_weight

    @property
    def realizations(self):
        """Flat list with all the realizations"""
        return sum(self.rlzs_by_smodel.values(), [])

    @property
    def weights(self):
        """Array with the weight of the realizations"""
        return numpy.array([rlz.weight for rlz in self.realizations])

    def combine_pmaps(self, pmap_by_grp):
        """
        :param pmap_by_grp: dictionary group string -> probability map
        :returns: a list of probability maps, one per realization
        """
        grp = list(pmap_by_grp)[0]  # pmap_by_grp must be non-empty
        num_levels = pmap_by_grp[grp].shape_y
        pmaps = [probability_map.ProbabilityMap(num_levels, 1)
                 for _ in self.realizations]
        array = self.by_grp()
        for grp in pmap_by_grp:
            for gsim_idx, rlzis in enumerate(array[grp]):
                pmap = pmap_by_grp[grp].extract(gsim_idx)
                for rlzi in rlzis:
                    pmaps[rlzi] |= pmap
        return pmaps

    def get_rlz(self, rlzstr):
        """
        Get a Realization instance for a string of the form 'rlz-\d+'
        """
        mo = re.match(r'rlz-(\d+)', rlzstr)
        if not mo:
            return
        return self.realizations[int(mo.group(1))]

    def _add_realizations(self, offset, lt_model, all_trts, gsim_rlzs):
        idx = numpy.arange(offset, offset + len(gsim_rlzs))
        rlzs = []
        for i, gsim_rlz in enumerate(gsim_rlzs):
            weight = float(lt_model.weight) * gsim_rlz.weight
            rlz = LtRealization(idx[i], lt_model.path, gsim_rlz, weight)
            self.gsim_by_trt.append(dict(zip(all_trts, gsim_rlz.value)))
            rlzs.append(rlz)
        self.rlzs_by_smodel[lt_model.ordinal] = rlzs

    def __repr__(self):
        pairs = []
        dic = self.by_grp()
        size = sum(len(dic[grp]) for grp in dic)
        for grp in sorted(dic):
            grp_id = int(grp[4:])
            gsims = self.csm_info.get_gsims(grp_id)
            for gsim_idx, rlzis in enumerate(dic[grp]):
                if len(rlzis) > 10:  # short representation
                    rlzis = ['%d realizations' % len(rlzis)]
                pairs.append(('%s,%s' % (grp_id, gsims[gsim_idx]), rlzis))
        return '<%s(size=%d, rlzs=%d)\n%s>' % (
            self.__class__.__name__, size, len(self.realizations),
            '\n'.join('%s: %s' % pair for pair in pairs))

def accept_path(path, ref_path):
    """
    :param path: a logic tree path (list or tuple of strings)
    :param ref_path: reference logic tree path
    :returns: True if `path` is consistent with `ref_path`, False otherwise

    >>> accept_path(['SM2'], ('SM2', 'a3b1'))
    False
    >>> accept_path(['SM2', '@'], ('SM2', 'a3b1'))
    True
    >>> accept_path(['@', 'a3b1'], ('SM2', 'a3b1'))
    True
    >>> accept_path('@@', ('SM2', 'a3b1'))
    True
    """
    if len(path) != len(ref_path):
        return False
    for a, b in zip(path, ref_path):
        if a != '@' and a != b:
            return False
    return True

def get_rlzs_assoc(cinfo, sm_lt_path=None, trts=None):
    """
    :param cinfo: a :class:`openquake.commonlib.source.CompositionInfo`
    :param sm_lt_path: logic tree path tuple used to select a source model
    :param trts: tectonic region types to accept
    """
    assoc = RlzsAssoc(cinfo)
    offset = 0
    trtset = set(cinfo.gsim_lt.all_trts)
    for smodel in cinfo.source_models:
        # discard source models with non-acceptable lt_path
        if sm_lt_path and not accept_path(smodel.path, sm_lt_path):
            continue

        # collect the effective tectonic region types and ruptures
        trts_ = set()
        for sg in smodel.src_groups:
            if sg.eff_ruptures:
                if (trts and sg.trt in trts) or not trts:
                    trts_.add(sg.trt)

        # recompute the GSIM logic tree if needed
        if trtset != trts_:
            before = cinfo.gsim_lt.get_num_paths()
            gsim_lt = cinfo.gsim_lt.reduce(trts_)
            after = gsim_lt.get_num_paths()
            if sm_lt_path and before > after:
                # print the warning only when saving the logic tree,
                # i.e. when called with sm_lt_path in store_rlz_info
                logging.warning('Reducing the logic tree of %s from %d to %d '
                                'realizations', smodel.name, before, after)
            gsim_rlzs = list(gsim_lt)
            all_trts = gsim_lt.all_trts
        else:
            gsim_rlzs = cinfo.gsim_rlzs
            all_trts = cinfo.gsim_lt.all_trts

        rlzs = cinfo._get_rlzs(smodel, gsim_rlzs, cinfo.seed + offset)
        assoc._add_realizations(offset, smodel, all_trts, rlzs)
        offset += len(rlzs)

    if assoc.realizations:
        assoc._init()
    return assoc
<|endoftext|>