An efficient harmonic oscillator term (Cython part) — MMTK User Guide 2.7.7 documentation

### Navigation

- index
- modules |
- MMTK User Guide 2.7.7 documentation »

# An efficient harmonic oscillator term (Cython part)

     1
     2
     3
     4
     5
     6
     7
     8
     9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    # Example for a forcefield implementation in Cython.

    #
    # Get all the required declarations
    #
    include "MMTK/python.pxi"
    include "MMTK/numeric.pxi"
    include "MMTK/core.pxi"
    include "MMTK/universe.pxi"
    include 'MMTK/forcefield.pxi'

    #
    # The force field term implementation.
    # The rules:
    #
    # - The class must inherit from EnergyTerm.
    #
    # - EnergyTerm.__init__() must be called with the arguments
    #   shown here. The third argument is the name of the EnergyTerm
    #   object, the fourth a tuple of the names of all the terms it
    #   implements (one object can implement several terms).
    #   The assignment to self.eval_func is essential, without it
    #   any energy evaluation will crash.
    #
    # - The function "evaluate" must have exactly the parameter
    #   list given in this example.
    #
    cdef class HarmonicOscillatorTerm(EnergyTerm):

        cdef int atom_index
        cdef double ref_x, ref_y, ref_z
        cdef double k

        def __init__(self, universe, atom_index, reference, force_constant):
            cdef ArrayType ref_array
            EnergyTerm.__init__(self, universe,
                                "harmonic_oscillator", ("harmonic_oscillator",))
            self.eval_func = <void *>HarmonicOscillatorTerm.evaluate
            self.atom_index = atom_index
            ref_array = reference.array
            self.ref_x = (<double *>ref_array.data)[0]
            self.ref_y = (<double *>ref_array.data)[1]
            self.ref_z = (<double *>ref_array.data)[2]
            self.k = force_constant

        # The function evaluate is called for every single energy
        # evaluation and should therefore be optimized for speed.
        # Its first argument is the global energy evaluator object,
        # which is needed only for parallelized energy terms.
        # The second argument is a C structure that contains all the
        # input data, in particular the particle configuration.
        # The third argument is a C structure that contains the
        # energy term fields and gradient arrays for storing the results.
        # For details, see MMTK_forcefield.pxi.
        cdef void evaluate(self, PyFFEvaluatorObject *eval,
                           energy_spec *input, energy_data *energy):
            cdef vector3 *coordinates, *gradients
            cdef double *fc
            cdef double dx, dy, dz
            cdef int n, offset
            coordinates = <vector3 *>input.coordinates.data
            dx = coordinates[self.atom_index][0] - self.ref_x
            dy = coordinates[self.atom_index][1] - self.ref_y
            dz = coordinates[self.atom_index][2] - self.ref_z
            energy.energy_terms[self.index] = 0.5*self.k*(dx*dx + dy*dy + dz*dz)
            energy.energy_terms[self.virial_index] -= self.k*(dx*dx + dy*dy + dz*dz)
            if energy.gradients != NULL:
                gradients = <vector3 *>(<PyArrayObject *> energy.gradients).data
                gradients[self.atom_index][0] += self.k*dx
                gradients[self.atom_index][1] += self.k*dy
                gradients[self.atom_index][2] += self.k*dz
            if energy.force_constants != NULL:
                fc = <double *>(<PyArrayObject *> energy.force_constants).data
                n = (<PyArrayObject *> energy.force_constants).dimensions[0]
                offset = (9*n+3)*self.atom_index
                fc[offset + 3*n*0 + 0] += self.k
                fc[offset + 3*n*1 + 1] += self.k
                fc[offset + 3*n*2 + 2] += self.k

### Navigation

- index
- modules |
- MMTK User Guide 2.7.7 documentation »