<commit_before># Copyright 2014 Sandra Giuliani
# drug_repo.py

# Mapping of known drugs from ChEMBL and DrugBank to their targets/domain
# architecture to identify suitable drug repositioning candidates for
# schistosomiasis.
# Please see README.md for more info.

############################################################################
### import modules, set up log and define constant variables
############################################################################

# import modules
import sys, re, string, os, fnmatch, shutil

# set up log
import logging
# set up log file to write to, it will be overwritten every time ('w' mode)
# leave this level setting to DEBUG
logging.basicConfig(filename='log_drug_repo.log', filemode='w',
                    level=logging.DEBUG)
logger = logging.getLogger(__name__)
# leave this level setting to DEBUG
logger.setLevel(logging.DEBUG)
# create console handler and set level to debug or other level
ch = logging.StreamHandler()
# CHANGE THIS TO TUNE LOGGING LEVEL from DEBUG/INFO/WARNING
ch.setLevel(logging.DEBUG)
# create formatter, you can add '%(levelname)s' for level name
# eg __main__ in the log
formatter = logging.Formatter('%(levelname)s: %(name)s - %(message)s')
# add formatter to ch
ch.setFormatter(formatter)
# add ch to logger
logger.addHandler(ch)

#logger.warning('warning_test')

# define CHEMBL_INPUT as the chembldrugs file
CHEMBL_INPUT = 'chembldrugs.txt'
# define CHEMBL_UNIPROT as the chemblID/uniprot mapping file
CHEMBL_UNIPROT = 'chembl_uniprot_mapping.txt'

############################################################################

############################################################################
### HEADER_TAB_COUNT HELPER FUNCTION
############################################################################
# find specific header in first line of tab-separated file
# and return column number of the header
def header_tab_count(tab_file, header):
  '''(str)->int
  Read first line of tab separated file tab_file and return column number
  of header
  >>>column_count("DEVELOPMENT_PHASE")
  3
  '''

  # read just first line of tab_file
  with open(tab_file, 'r') as f:
    first = f.readline()
  # set counter to 0
  col_number = 0
  # loop until word matches the header
  while not first.split("\t")[col_number] == header:
    col_number = col_number + 1
  #debug check
  logger.debug('the header tab count thing works')
  # return column number
  return col_number

############################################################################

############################################################################
### FILE_TO_LINES HELPER FUNCTION
############################################################################
# read file using readlines approach and return the lines
def file_to_lines(text_file):
  input_file = open(text_file, 'r')
  lines = input_file.readlines()
  input_file.close()
  #logger.debug(lines)
  return lines

############################################################################

############################################################################
### SWAP_DIC HELPER FUNCTION
############################################################################
# read tab-separated mapping file with header and return dictionary with
# second column as key and first column as values - created for the
# chemblID uniprot mapping file
def swap_dic(tab_file):
  lines = file_to_lines(tab_file)
  swap_dictionary = {}
  counter = 0
  # iterate over lines
  for i in range(1,len(lines)):
    # split tab
    splitline = lines[i].split("\t")
    counter = counter + 1
    #logger.debug(splitline[1])
    # create dictionary, stripping the carriage return
    swap_dictionary[splitline[1].rstrip('\r\n')] = (splitline[0])
  #logger.debug(swap_dictionary)
  return swap_dictionary

############################################################################

############################################################################
### PROCESS_CHEMBL FUNCTION
############################################################################

# in the end we need a dictionary of chembl ids vs uniprot ids
def process_chembl():
  '''read chembl drug input file, filter, get various info and return
  dictionary of '''
  # open chembldrugs.txt for reading
  lines = file_to_lines(CHEMBL_INPUT)
  logger.info('The number of drugs listed in the input file is '
              + str(len(lines)-1))
  # get the headers
  headers = lines[0]
  # remove duplicate spaces
  headersnospace = " ".join(headers.split())
  # print headers list in lowercase
  # print('The headers are: '+headersnospace.lower())

  col_phase = header_tab_count(CHEMBL_INPUT,"DEVELOPMENT_PHASE")
  col_type = header_tab_count(CHEMBL_INPUT,"DRUG_TYPE")
  col_chemblid = header_tab_count(CHEMBL_INPUT,"CHEMBL_ID")

  logger.info('The column with the development_phase info is the ' +
              str(col_phase+1) + 'th; the one with the drug_type info is the '
              + str(col_type+1) + 'th.')

  # count the drugs in the 4+1 classes of clinical phase
  # set counters for clinical phases to zero
  phase1 = 0
  phase2 = 0
  phase3 = 0
  phase4 = 0
  phase_unknown =0
  # iterate over rows, excluding the header row
  for i in range(1,len(lines)):
    # tab separate each row
    rowsplit = lines[i].split("\t")
    if (rowsplit[col_phase] == '1'):
      phase1 = phase1 + 1
    elif (rowsplit[col_phase] == '2'):
      phase2 = phase2 + 1
    elif (rowsplit[col_phase] == '3'):
      phase3 = phase3 + 1
    elif (rowsplit[col_phase] == '4'):
      phase4 = phase4 + 1
    else:
      phase_unknown = phase_unknown + 1

  logger.info('Number of drugs in phase 1 is: ' + str(phase1) +
              '; in phase 2 is: ' + str(phase2) + '; in phase 3 is: ' +
              str(phase3)+'; in phase 4 is: ' + str(phase4) +
              '; in unknown phase is: ' + str(phase_unknown) + '.')

  # open the file to write to
  stripped = open('chembldrugs_stripped.txt', 'w')
  # set counter to zero, this is just to know how many lines we end up writing
  total_stripped_lines = 0
  # look over lines, excluding header
  for y in range(1,len(lines)):
    # tab separate each row
    rowsplit2 = lines[y].split("\t")
    # check if they are phase 4 or unknown
    # can also add phase 3 from here rowsplit2[col_phase] == '3')
    if (rowsplit2[col_phase] == '4') or (rowsplit2[col_phase] == ''):
      # increase the useless counter
      total_stripped_lines = total_stripped_lines + 1
      # write to the file the stripped lines
      stripped.write(lines[y])
      # print(rowsplit2[column])
  # print friendly statement
  logger.info('We have written to the file chembl_stripped.txt' +
      'only the entries in phase 4 or with unkown phase' +
      ', for a total number of '+ str(total_stripped_lines)+' drugs.')
  # close the file we wrote to
  stripped.close()

  # we open the stripped file for reading
  stripped2 = open('chembldrugs_stripped.txt', 'r')
  # reading lines
  lines2 = stripped2.readlines()
  # closing the stripped file
  stripped2.close()
  filtered = open('chembldrugs_filt.txt', 'w')
  small_mol_count = 0
  # look over, note here there is no header
  for x in range(len(lines2)):
    # tab separate
    rowsplit3 = lines2[x].split("\t")
    # drug_type = ''
    if rowsplit3[col_type] == 'Synthetic Small Molecule':
      filtered.write(lines2[x])
      small_mol_count = small_mol_count +1
  #chembl_dic[lines2[col_chemblid]] = [lines2[]]
  logger.debug(small_mol_count)
  # create dictionary from the chembl/uniprot mapping file
  chembl_uniprot_map_dic = swap_dic(CHEMBL_UNIPROT)
  # empty dictionary on which to store filter values
  chembldrugs_uniprot_dic = {}

############################################################################

############################################################################
### PROCESS_DRUGBANK FUNCTION
############################################################################

# process drugbank file to produce a dictionary of drugs vs uniprot ids
#def process_drugbank():
  # implement function

############################################################################

############################################################################
### MAIN FUNCTION
############################################################################

# call process_chembl,
# TO ADD: call process_drugbank, merge uniprot codes

def main():
  process_chembl()

############################################################################

############################################################################
### call main function, prevent excecution on import
############################################################################
if __name__ == "__main__":
  main()
############################################################################
<commit_msg>minor changes drug_repo<commit_after># Copyright 2014 Sandra Giuliani
# drug_repo.py

# Mapping of known drugs from ChEMBL and DrugBank to their targets/domain
# architecture to identify suitable drug repositioning candidates for
# schistosomiasis.
# Please see README.md for more info.

############################################################################
### import modules, set up log and define constant variables
############################################################################

# import modules
import sys, re, string, os, fnmatch, shutil

# set up log
import logging
# set up log file to write to, it will be overwritten every time ('w' mode)
# leave this level setting to DEBUG
logging.basicConfig(filename='log_drug_repo.log', filemode='w',
                    level=logging.DEBUG)
logger = logging.getLogger(__name__)
# leave this level setting to DEBUG
logger.setLevel(logging.DEBUG)
# create console handler
ch = logging.StreamHandler()
# CHANGE THIS TO TUNE LOGGING LEVEL from DEBUG/INFO/WARNING
ch.setLevel(logging.DEBUG)
# create formatter, you can add '%(levelname)s' for level name
# eg __main__ in the log
formatter = logging.Formatter('%(levelname)s: %(name)s - %(message)s')
# add formatter to ch
ch.setFormatter(formatter)
# add ch to logger
logger.addHandler(ch)

#logger.warning('warning_test')

# define CHEMBL_INPUT as the chembldrugs file
CHEMBL_INPUT = 'chembldrugs.txt'
# define CHEMBL_UNIPROT as the chemblID/uniprot mapping file
CHEMBL_UNIPROT = 'chembl_uniprot_mapping.txt'
############################################################################

############################################################################
### HEADER_TAB_COUNT HELPER FUNCTION
############################################################################
# find specific header in first line of tab-separated file
# and return column number of the header
def header_tab_count(tab_file, header):
  '''(str)->int
  Read first line of tab separated file tab_file and return column number
  of header
  >>>column_count("DEVELOPMENT_PHASE")
  3
  '''

  # read just first line of tab_file
  with open(tab_file, 'r') as f:
    first = f.readline()
  # set counter to 0
  col_number = 0
  # loop until word matches the header
  while not first.split("\t")[col_number] == header:
    col_number = col_number + 1
  #debug check
  logger.debug('the header tab count thing works')
  # return column number
  return col_number
############################################################################

############################################################################
### FILE_TO_LINES HELPER FUNCTION
############################################################################
# read file using readlines approach and return the lines
def file_to_lines(text_file):
  input_file = open(text_file, 'r')
  lines = input_file.readlines()
  input_file.close()
  #logger.debug(lines)
  return lines
############################################################################

############################################################################
### SWAP_DIC HELPER FUNCTION
############################################################################
# read tab-separated mapping file with header and return dictionary with
# second column as key and first column as values - created for the
# chemblID uniprot mapping file
def swap_dic(tab_file):
  lines = file_to_lines(tab_file)
  swap_dictionary = {}
  counter = 0
  # iterate over lines
  for i in range(1,len(lines)):
    # split tab
    splitline = lines[i].split("\t")
    counter = counter + 1
    #logger.debug(splitline[1])
    # create dictionary, stripping the carriage return
    swap_dictionary[splitline[1].rstrip('\r\n')] = (splitline[0])
  #logger.debug(swap_dictionary)
  return swap_dictionary
############################################################################

############################################################################
### PROCESS_CHEMBL FUNCTION
############################################################################

# in the end we need a dictionary of chembl ids vs uniprot ids
def process_chembl():
  '''read chembl drug input file, filter, get various info and return
  dictionary of '''
  # open chembldrugs.txt for reading
  lines = file_to_lines(CHEMBL_INPUT)
  logger.info('The number of drugs listed in the input file is '
              + str(len(lines)-1))
  # get the headers
  headers = lines[0]
  # remove duplicate spaces
  headersnospace = " ".join(headers.split())
  # print headers list in lowercase
  # print('The headers are: '+headersnospace.lower())

  col_phase = header_tab_count(CHEMBL_INPUT,"DEVELOPMENT_PHASE")
  col_type = header_tab_count(CHEMBL_INPUT,"DRUG_TYPE")
  col_chemblid = header_tab_count(CHEMBL_INPUT,"CHEMBL_ID")

  logger.info('The column with the development_phase info is the ' +
              str(col_phase+1) + 'th; the one with the drug_type info is the '
              + str(col_type+1) + 'th.')

  # count the drugs in the 4+1 classes of clinical phase
  # set counters for clinical phases to zero
  phase1 = 0
  phase2 = 0
  phase3 = 0
  phase4 = 0
  phase_unknown =0
  # iterate over rows, excluding the header row
  for i in range(1,len(lines)):
    # tab separate each row
    rowsplit = lines[i].split("\t")
    if (rowsplit[col_phase] == '1'):
      phase1 = phase1 + 1
    elif (rowsplit[col_phase] == '2'):
      phase2 = phase2 + 1
    elif (rowsplit[col_phase] == '3'):
      phase3 = phase3 + 1
    elif (rowsplit[col_phase] == '4'):
      phase4 = phase4 + 1
    else:
      phase_unknown = phase_unknown + 1

  logger.info('Number of drugs in phase 1 is: ' + str(phase1) +
              '; in phase 2 is: ' + str(phase2) + '; in phase 3 is: ' +
              str(phase3)+'; in phase 4 is: ' + str(phase4) +
              '; in unknown phase is: ' + str(phase_unknown) + '.')

  # open the file to write to
  stripped = open('chembldrugs_stripped.txt', 'w')
  # set counter to zero, this is just to know how many lines we end up writing
  total_stripped_lines = 0
  # look over lines, excluding header
  for y in range(1,len(lines)):
    # tab separate each row
    rowsplit2 = lines[y].split("\t")
    # check if they are phase 4 or unknown
    # can also add phase 3 from here rowsplit2[col_phase] == '3')
    if (rowsplit2[col_phase] == '4') or (rowsplit2[col_phase] == ''):
      # increase the useless counter
      total_stripped_lines = total_stripped_lines + 1
      # write to the file the stripped lines
      stripped.write(lines[y])
      # print(rowsplit2[column])
  # print friendly statement
  logger.info('We have written to the file chembl_stripped.txt' +
      'only the entries in phase 4 or with unkown phase' +
      ', for a total number of '+ str(total_stripped_lines)+' drugs.')
  # close the file we wrote to
  stripped.close()

  # we open the stripped file for reading
  stripped2 = open('chembldrugs_stripped.txt', 'r')
  # reading lines
  lines2 = stripped2.readlines()
  # closing the stripped file
  stripped2.close()
  filtered = open('chembldrugs_filt.txt', 'w')
  small_mol_count = 0
  # look over, note here there is no header
  for x in range(len(lines2)):
    # tab separate
    rowsplit3 = lines2[x].split("\t")
    # drug_type = ''
    if rowsplit3[col_type] == 'Synthetic Small Molecule':
      filtered.write(lines2[x])
      small_mol_count = small_mol_count +1
  #chembl_dic[lines2[col_chemblid]] = [lines2[]]
  logger.debug(small_mol_count)
  # create dictionary from the chembl/uniprot mapping file
  chembl_uniprot_map_dic = swap_dic(CHEMBL_UNIPROT)
  # empty dictionary on which to store filtered values
  chembldrugs_uniprot_dic = {}

############################################################################

############################################################################
### PROCESS_DRUGBANK FUNCTION
############################################################################

# process drugbank file to produce a dictionary of drugs vs uniprot ids
#def process_drugbank():
  # implement function

############################################################################

############################################################################
### MAIN FUNCTION
############################################################################

# call process_chembl,
# TO ADD: call process_drugbank, merge uniprot codes

def main():
  process_chembl()

############################################################################

############################################################################
### call main function, prevent excecution on import
############################################################################
if __name__ == "__main__":
  main()
############################################################################
<|endoftext|>


<commit_before># Copyright (c) 2009-2019 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause License.

# Maintainer: joaander / All Developers are free to add commands for new features

R""" Bond potentials.

Bonds add forces between specified pairs of particles and are typically used to
model chemical bonds between two particles.

By themselves, bonds that have been specified in an initial configuration do nothing. Only when you
specify an bond force (i.e. bond.harmonic), are forces actually calculated between the
listed particles.
"""

from hoomd import _hoomd
from hoomd.md import _md
from hoomd.md.force import _Force
from hoomd.md import force
from hoomd.typeparam import TypeParameter
from hoomd.parameterdicts import TypeParameterDict
import hoomd

import math

class _Bond(_Force):
    """Constructs the bond potential.

    A bond in hoomd reflects a PotentialBond in c++. It is responsible for all
    high-level management that happens behind the scenes for hoomd writers.
    1) The instance of the c++ bond force itself is tracked and added to the
    System
    2) methods are provided for disabling the force from being added to the net
    force on each particle
    """
    def attach(self, simulation):
        """Create the c++ mirror class."""
        if not simulation.device.cpp_exec_conf.isCUDAEnabled():
            cpp_cls = getattr(_md, self._cpp_class_name)
        else:
            cpp_cls = getattr(_md, self._cpp_class_name + "GPU")

        # TODO remove string argument
        self._cpp_obj = cpp_cls(simulation.state._cpp_sys_def, "")

        super().attach(simulation)

class Harmonic(_Bond):
    R""" Harmonic bond potential.

    Args:
        name (str): Name of the bond instance.

    :py:class:`Harmonic` specifies a harmonic potential energy between the two
    particles in each defined bond.

    .. math::

        V(r) = \frac{1}{2} k \left( r - r_0 \right)^2

    where :math:`\vec{r}` is the vector pointing from one particle to the other
    in the bond.

    Coefficients:

    - :math:`k` - force constant ``k`` (in units of energy/distance^2)
    - :math:`r_0` - bond rest length ``r0`` (in distance units)
    """
    _cpp_class_name = "PotentialBondHarmonic"
    def __init__(self):
        params = TypeParameter("params", "bond_types",
                               TypeParameterDict(k=float, r0=float, len_keys=1)
                               )
        self._add_typeparam(params)

class FENE(_Bond):
    R""" FENE bond potential.

    Args:
        name (str): Name of the bond instance.

    :py:class:`FENE` specifies a FENE potential energy between the two particles
    in each defined bond.

    .. math::

        V(r) = - \frac{1}{2} k r_0^2 \ln \left( 1 - \left( \frac{r -
               \Delta}{r_0} \right)^2 \right) + V_{\mathrm{WCA}}(r)

    where :math:`\vec{r}` is the vector pointing from one particle to the other
    in the bond, :math:`\Delta = (d_i + d_j)/2 - 1`, :math:`d_i` is the diameter
    of particle :math:`i`, and

    .. math::
        :nowrap:

        \begin{eqnarray*}
        V_{\mathrm{WCA}}(r)  = & 4 \varepsilon \left[ \left( \frac{\sigma}{r -
                                 \Delta} \right)^{12} - \left( \frac{\sigma}{r -
                                 \Delta} \right)^{6} \right]  + \varepsilon
                               & r-\Delta < 2^{\frac{1}{6}}\sigma\\
                             = & 0
                               & r-\Delta \ge 2^{\frac{1}{6}}\sigma
        \end{eqnarray*}

    Coefficients:

    - :math:`k` - attractive force strength ``k`` (in units of
        energy/distance^2)
    - :math:`r_0` - size parameter ``r0`` (in distance units)
    - :math:`\varepsilon` - repulsive force strength ``epsilon`` (in energy
        units)
    - :math:`\sigma` - repulsive force interaction distance ``sigma`` (in
        distance units)
    """
    _cpp_class_name = "PotentialBondFENE"

    def __init__(self):
        params = TypeParameter("params", "bond_types",
                               TypeParameterDict(k=float,
                                                 r0=float,
                                                 epsilon=float,
                                                 sigma=float,
                                                 len_keys=1)
                               )
        self._add_typeparam(params)

def _table_eval(r, rmin, rmax, V, F, width):
      dr = (rmax - rmin) / float(width-1);
      i = int(round((r - rmin)/dr))
      return (V[i], F[i])

class table(force._force):
    R""" Tabulated bond potential.

    Args:
        width (int): Number of points to use to interpolate V and F
        name (str): Name of the potential instance

    :py:class:`table` specifies that a tabulated bond potential should be applied between the two particles in each
    defined bond.

    The force :math:`\vec{F}` is (in force units) and the potential :math:`V(r)` is (in energy units):

    .. math::
        :nowrap:

        \begin{eqnarray*}
        \vec{F}(\vec{r})     = & 0                           & r < r_{\mathrm{min}} \\
                             = & F_{\mathrm{user}}(r)\hat{r} & r < r_{\mathrm{max}} \\
                             = & 0                           & r \ge r_{\mathrm{max}} \\
                             \\
        V(r)       = & 0                    & r < r_{\mathrm{min}} \\
                   = & V_{\mathrm{user}}(r) & r < r_{\mathrm{max}} \\
                   = & 0                    & r \ge r_{\mathrm{max}} \\
        \end{eqnarray*}

    where :math:`\vec{r}` is the vector pointing from one particle to the other in the bond.  Care should be taken to
    define the range of the bond so that it is not possible for the distance between two bonded particles to be outside the
    specified range.  On the CPU, this will throw an error.  On the GPU, this will throw an error if GPU error checking is enabled.

    :math:`F_{\mathrm{user}}(r)` and :math:`V_{\mathrm{user}}(r)` are evaluated on *width* grid points between
    :math:`r_{\mathrm{min}}` and :math:`r_{\mathrm{max}}`. Values are interpolated linearly between grid points.
    For correctness, you must specify the force defined by: :math:`F = -\frac{\partial V}{\partial r}`

    The following coefficients must be set for each bond type:

    - :math:`F_{\mathrm{user}}(r)` and :math:`V_{\mathrm{user}}(r)` - evaluated by ``func`` (see example)
    - coefficients passed to ``func`` - ``coeff`` (see example)
    - :math:`r_{\mathrm{min}}` - ``rmin`` (in distance units)
    - :math:`r_{\mathrm{max}}` - ``rmax`` (in distance units)

    The table *width* is set once when bond.table is specified.
    There are two ways to specify the other parameters.

    .. rubric:: Set table from a given function

    When you have a functional form for V and F, you can enter that
    directly into python. :py:class:`table` will evaluate the given function over *width* points between *rmin* and *rmax*
    and use the resulting values in the table::

        def harmonic(r, rmin, rmax, kappa, r0):
           V = 0.5 * kappa * (r-r0)**2;
           F = -kappa*(r-r0);
           return (V, F)

        btable = bond.table(width=1000)
        btable.bond_coeff.set('bond1', func=harmonic, rmin=0.2, rmax=5.0, coeff=dict(kappa=330, r0=0.84))
        btable.bond_coeff.set('bond2', func=harmonic, rmin=0.2, rmax=5.0, coeff=dict(kappa=30, r0=1.0))

    .. rubric:: Set a table from a file

    When you have no function for for *V* or *F*, or you otherwise have the data listed in a file, :py:class:`table` can use the given
    values directly. You must first specify the number of rows in your tables when initializing bond.table. Then use
    :py:meth:`set_from_file()` to read the file::

        btable = bond.table(width=1000)
        btable.set_from_file('polymer', 'btable.file')

    Note:
        For potentials that diverge near r=0, make sure to set ``rmin`` to a reasonable value. If a potential does
        not diverge near r=0, then a setting of ``rmin=0`` is valid.

    Note:
        Ensure that ``rmin`` and ``rmax`` cover the range of possible bond lengths. When gpu error checking is on, a error will
        be thrown if a bond distance is outside than this range.
    """
    def __init__(self, width, name=None):

        # initialize the base class
        force._force.__init__(self, name);

        # create the c++ mirror class
        if not hoomd.context.current.device.cpp_exec_conf.isCUDAEnabled():
            self.cpp_force = _md.BondTablePotential(hoomd.context.current.system_definition, int(width), self.name);
        else:
            self.cpp_force = _md.BondTablePotentialGPU(hoomd.context.current.system_definition, int(width), self.name);

        hoomd.context.current.system.addCompute(self.cpp_force, self.force_name);

        # setup the coefficients matrix
        self.bond_coeff = coeff();

        # stash the width for later use
        self.width = width;

    def update_bond_table(self, btype, func, rmin, rmax, coeff):
        # allocate arrays to store V and F
        Vtable = _hoomd.std_vector_scalar();
        Ftable = _hoomd.std_vector_scalar();

        # calculate dr
        dr = (rmax - rmin) / float(self.width-1);

        # evaluate each point of the function
        for i in range(0, self.width):
            r = rmin + dr * i;
            (V,F) = func(r, rmin, rmax, **coeff);

            # fill out the tables
            Vtable.append(V);
            Ftable.append(F);

        # pass the tables on to the underlying cpp compute
        self.cpp_force.setTable(btype, Vtable, Ftable, rmin, rmax);

    def update_coeffs(self):
        # check that the bond coefficients are valid
        if not self.bond_coeff.verify(["func", "rmin", "rmax", "coeff"]):
            hoomd.context.current.device.cpp_msg.error("Not all bond coefficients are set for bond.table\n");
            raise RuntimeError("Error updating bond coefficients");

        # set all the params
        ntypes = hoomd.context.current.system_definition.getBondData().getNTypes();
        type_list = [];
        for i in range(0,ntypes):
            type_list.append(hoomd.context.current.system_definition.getBondData().getNameByType(i));

        # loop through all of the unique type bonds and evaluate the table
        for i in range(0,ntypes):
            func = self.bond_coeff.get(type_list[i], "func");
            rmin = self.bond_coeff.get(type_list[i], "rmin");
            rmax = self.bond_coeff.get(type_list[i], "rmax");
            coeff = self.bond_coeff.get(type_list[i], "coeff");

            self.update_bond_table(i, func, rmin, rmax, coeff);

    def set_from_file(self, bondname, filename):
        R""" Set a bond pair interaction from a file.

        Args:
            bondname (str): Name of bond
            filename (str): Name of the file to read

        The provided file specifies V and F at equally spaced r values.
        Example::

            #r  V    F
            1.0 2.0 -3.0
            1.1 3.0 -4.0
            1.2 2.0 -3.0
            1.3 1.0 -2.0
            1.4 0.0 -1.0
            1.5 -1.0 0.0

        The first r value sets ``rmin``, the last sets ``rmax``. Any line with # as the first non-whitespace character is
        is treated as a comment. The ``r`` values must monotonically increase and be equally spaced. The table is read
        directly into the grid points used to evaluate :math:`F_{\mathrm{user}}(r)` and :math:`V_{\mathrm{user}}(r)`.
        """

        # open the file
        f = open(filename);

        r_table = [];
        V_table = [];
        F_table = [];

        # read in lines from the file
        for line in f.readlines():
            line = line.strip();

            # skip comment lines
            if line[0] == '#':
                continue;

            # split out the columns
            cols = line.split();
            values = [float(f) for f in cols];

            # validate the input
            if len(values) != 3:
                hoomd.context.current.device.cpp_msg.error("bond.table: file must have exactly 3 columns\n");
                raise RuntimeError("Error reading table file");

            # append to the tables
            r_table.append(values[0]);
            V_table.append(values[1]);
            F_table.append(values[2]);

        # validate input
        if self.width != len(r_table):
            hoomd.context.current.device.cpp_msg.error("bond.table: file must have exactly " + str(self.width) + " rows\n");
            raise RuntimeError("Error reading table file");

        # extract rmin and rmax
        rmin_table = r_table[0];
        rmax_table = r_table[-1];

        # check for even spacing
        dr = (rmax_table - rmin_table) / float(self.width-1);
        for i in range(0,self.width):
            r = rmin_table + dr * i;
            if math.fabs(r - r_table[i]) > 1e-3:
                hoomd.context.current.device.cpp_msg.error("bond.table: r must be monotonically increasing and evenly spaced\n");
                raise RuntimeError("Error reading table file");

        self.bond_coeff.set(bondname, func=_table_eval, rmin=rmin_table, rmax=rmax_table, coeff=dict(V=V_table, F=F_table, width=self.width))
<commit_msg>updated docstrings for bond forces (except table)<commit_after># Copyright (c) 2009-2019 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause License.

# Maintainer: joaander / All Developers are free to add commands for new features

R""" Bond potentials.

Bonds add forces between specified pairs of particles and are typically used to
model chemical bonds between two particles.

By themselves, bonds that have been specified in an initial configuration do nothing. Only when you
specify an bond force (i.e. bond.harmonic), are forces actually calculated between the
listed particles.
"""

from hoomd import _hoomd
from hoomd.md import _md
from hoomd.md.force import _Force
from hoomd.md import force
from hoomd.typeparam import TypeParameter
from hoomd.parameterdicts import TypeParameterDict
import hoomd

import math

class _Bond(_Force):
    """Constructs the bond potential.

    A bond in hoomd reflects a PotentialBond in c++. It is responsible for all
    high-level management that happens behind the scenes for hoomd writers.
    1) The instance of the c++ bond force itself is tracked and added to the
    System
    2) methods are provided for disabling the force from being added to the net
    force on each particle
    """
    def attach(self, simulation):
        """Create the c++ mirror class."""
        if not simulation.device.cpp_exec_conf.isCUDAEnabled():
            cpp_cls = getattr(_md, self._cpp_class_name)
        else:
            cpp_cls = getattr(_md, self._cpp_class_name + "GPU")

        # TODO remove string argument
        self._cpp_obj = cpp_cls(simulation.state._cpp_sys_def, "")

        super().attach(simulation)

class Harmonic(_Bond):
    R""" Harmonic bond potential.

    :py:class:`Harmonic` specifies a harmonic potential energy between the two
    particles in each defined bond.

    .. math::

        V(r) = \frac{1}{2} k \left( r - r_0 \right)^2

    where :math:`\vec{r}` is the vector pointing from one particle to the other
    in the bond.

    Attributes:
        params (TypeParameter[``bond type``, dict]):
            The parameter of the harmonic bonds for each particle type.
            The dictionary has the following keys:

            * ``k`` (`float`, **required**) - potential constant
              (in units of energy/distance^2)

            * ``r0`` (`float`, **required**) - rest length
              (in units distance)

    Examples::

        harmonic = bond.Harmonic()
        harmonic.params['polymer'] = dict(k=3.0, r0=2.38)
        harmonic.params['backbone'] = dict(k=10.0, r0=1.0)
    """
    _cpp_class_name = "PotentialBondHarmonic"
    def __init__(self):
        params = TypeParameter("params", "bond_types",
                               TypeParameterDict(k=float, r0=float, len_keys=1)
                               )
        self._add_typeparam(params)

class FENE(_Bond):
    R""" FENE bond potential.

    :py:class:`FENE` specifies a FENE potential energy between the two particles
    in each defined bond.

    .. math::

        V(r) = - \frac{1}{2} k r_0^2 \ln \left( 1 - \left( \frac{r -
               \Delta}{r_0} \right)^2 \right) + V_{\mathrm{WCA}}(r)

    where :math:`\vec{r}` is the vector pointing from one particle to the other
    in the bond, :math:`\Delta = (d_i + d_j)/2 - 1`, :math:`d_i` is the diameter
    of particle :math:`i`, and

    .. math::
        :nowrap:

        \begin{eqnarray*}
        V_{\mathrm{WCA}}(r)  = & 4 \varepsilon \left[ \left( \frac{\sigma}{r -
                                 \Delta} \right)^{12} - \left( \frac{\sigma}{r -
                                 \Delta} \right)^{6} \right]  + \varepsilon
                               & r-\Delta < 2^{\frac{1}{6}}\sigma\\
                             = & 0
                               & r-\Delta \ge 2^{\frac{1}{6}}\sigma
        \end{eqnarray*}

    Attributes:
        params (TypeParameter[``bond type``, dict]):
            The parameter of the FENE potential bonds.
            The dictionary has the following keys:

            * ``k`` (`float`, **required**) - attractive force strength
              (in units of energy/distance^2)

            * ``r0`` (`float`, **required**) - size parameter
              (in units distance)

            * ``epsilon`` (`float`, **required**) - repulsive force strength
              (in units of energy)

            * ``sigma`` (`float`, **required**) - repulsive force interaction distance
              (in units of distance)

    Examples::

        bond_potential = bond.FENE()
        bond_potential.params['molecule'] = dict(k=3.0, r0=2.38, epsilon=1.0, sigma=1.0)
        bond_potential.params['backbone'] = dict(k=10.0, r0=1.0, epsilon=0.8, sigma=1.2)

    """
    _cpp_class_name = "PotentialBondFENE"

    def __init__(self):
        params = TypeParameter("params", "bond_types",
                               TypeParameterDict(k=float,
                                                 r0=float,
                                                 epsilon=float,
                                                 sigma=float,
                                                 len_keys=1)
                               )
        self._add_typeparam(params)

def _table_eval(r, rmin, rmax, V, F, width):
      dr = (rmax - rmin) / float(width-1);
      i = int(round((r - rmin)/dr))
      return (V[i], F[i])

class table(force._force):
    R""" Tabulated bond potential.

    Args:
        width (int): Number of points to use to interpolate V and F
        name (str): Name of the potential instance

    :py:class:`table` specifies that a tabulated bond potential should be applied between the two particles in each
    defined bond.

    The force :math:`\vec{F}` is (in force units) and the potential :math:`V(r)` is (in energy units):

    .. math::
        :nowrap:

        \begin{eqnarray*}
        \vec{F}(\vec{r})     = & 0                           & r < r_{\mathrm{min}} \\
                             = & F_{\mathrm{user}}(r)\hat{r} & r < r_{\mathrm{max}} \\
                             = & 0                           & r \ge r_{\mathrm{max}} \\
                             \\
        V(r)       = & 0                    & r < r_{\mathrm{min}} \\
                   = & V_{\mathrm{user}}(r) & r < r_{\mathrm{max}} \\
                   = & 0                    & r \ge r_{\mathrm{max}} \\
        \end{eqnarray*}

    where :math:`\vec{r}` is the vector pointing from one particle to the other in the bond.  Care should be taken to
    define the range of the bond so that it is not possible for the distance between two bonded particles to be outside the
    specified range.  On the CPU, this will throw an error.  On the GPU, this will throw an error if GPU error checking is enabled.

    :math:`F_{\mathrm{user}}(r)` and :math:`V_{\mathrm{user}}(r)` are evaluated on *width* grid points between
    :math:`r_{\mathrm{min}}` and :math:`r_{\mathrm{max}}`. Values are interpolated linearly between grid points.
    For correctness, you must specify the force defined by: :math:`F = -\frac{\partial V}{\partial r}`

    The following coefficients must be set for each bond type:

    - :math:`F_{\mathrm{user}}(r)` and :math:`V_{\mathrm{user}}(r)` - evaluated by ``func`` (see example)
    - coefficients passed to ``func`` - ``coeff`` (see example)
    - :math:`r_{\mathrm{min}}` - ``rmin`` (in distance units)
    - :math:`r_{\mathrm{max}}` - ``rmax`` (in distance units)

    The table *width* is set once when bond.table is specified.
    There are two ways to specify the other parameters.

    .. rubric:: Set table from a given function

    When you have a functional form for V and F, you can enter that
    directly into python. :py:class:`table` will evaluate the given function over *width* points between *rmin* and *rmax*
    and use the resulting values in the table::

        def harmonic(r, rmin, rmax, kappa, r0):
           V = 0.5 * kappa * (r-r0)**2;
           F = -kappa*(r-r0);
           return (V, F)

        btable = bond.table(width=1000)
        btable.bond_coeff.set('bond1', func=harmonic, rmin=0.2, rmax=5.0, coeff=dict(kappa=330, r0=0.84))
        btable.bond_coeff.set('bond2', func=harmonic, rmin=0.2, rmax=5.0, coeff=dict(kappa=30, r0=1.0))

    .. rubric:: Set a table from a file

    When you have no function for for *V* or *F*, or you otherwise have the data listed in a file, :py:class:`table` can use the given
    values directly. You must first specify the number of rows in your tables when initializing bond.table. Then use
    :py:meth:`set_from_file()` to read the file::

        btable = bond.table(width=1000)
        btable.set_from_file('polymer', 'btable.file')

    Note:
        For potentials that diverge near r=0, make sure to set ``rmin`` to a reasonable value. If a potential does
        not diverge near r=0, then a setting of ``rmin=0`` is valid.

    Note:
        Ensure that ``rmin`` and ``rmax`` cover the range of possible bond lengths. When gpu error checking is on, a error will
        be thrown if a bond distance is outside than this range.
    """
    def __init__(self, width, name=None):

        # initialize the base class
        force._force.__init__(self, name);

        # create the c++ mirror class
        if not hoomd.context.current.device.cpp_exec_conf.isCUDAEnabled():
            self.cpp_force = _md.BondTablePotential(hoomd.context.current.system_definition, int(width), self.name);
        else:
            self.cpp_force = _md.BondTablePotentialGPU(hoomd.context.current.system_definition, int(width), self.name);

        hoomd.context.current.system.addCompute(self.cpp_force, self.force_name);

        # setup the coefficients matrix
        self.bond_coeff = coeff();

        # stash the width for later use
        self.width = width;

    def update_bond_table(self, btype, func, rmin, rmax, coeff):
        # allocate arrays to store V and F
        Vtable = _hoomd.std_vector_scalar();
        Ftable = _hoomd.std_vector_scalar();

        # calculate dr
        dr = (rmax - rmin) / float(self.width-1);

        # evaluate each point of the function
        for i in range(0, self.width):
            r = rmin + dr * i;
            (V,F) = func(r, rmin, rmax, **coeff);

            # fill out the tables
            Vtable.append(V);
            Ftable.append(F);

        # pass the tables on to the underlying cpp compute
        self.cpp_force.setTable(btype, Vtable, Ftable, rmin, rmax);

    def update_coeffs(self):
        # check that the bond coefficients are valid
        if not self.bond_coeff.verify(["func", "rmin", "rmax", "coeff"]):
            hoomd.context.current.device.cpp_msg.error("Not all bond coefficients are set for bond.table\n");
            raise RuntimeError("Error updating bond coefficients");

        # set all the params
        ntypes = hoomd.context.current.system_definition.getBondData().getNTypes();
        type_list = [];
        for i in range(0,ntypes):
            type_list.append(hoomd.context.current.system_definition.getBondData().getNameByType(i));

        # loop through all of the unique type bonds and evaluate the table
        for i in range(0,ntypes):
            func = self.bond_coeff.get(type_list[i], "func");
            rmin = self.bond_coeff.get(type_list[i], "rmin");
            rmax = self.bond_coeff.get(type_list[i], "rmax");
            coeff = self.bond_coeff.get(type_list[i], "coeff");

            self.update_bond_table(i, func, rmin, rmax, coeff);

    def set_from_file(self, bondname, filename):
        R""" Set a bond pair interaction from a file.

        Args:
            bondname (str): Name of bond
            filename (str): Name of the file to read

        The provided file specifies V and F at equally spaced r values.
        Example::

            #r  V    F
            1.0 2.0 -3.0
            1.1 3.0 -4.0
            1.2 2.0 -3.0
            1.3 1.0 -2.0
            1.4 0.0 -1.0
            1.5 -1.0 0.0

        The first r value sets ``rmin``, the last sets ``rmax``. Any line with # as the first non-whitespace character is
        is treated as a comment. The ``r`` values must monotonically increase and be equally spaced. The table is read
        directly into the grid points used to evaluate :math:`F_{\mathrm{user}}(r)` and :math:`V_{\mathrm{user}}(r)`.
        """

        # open the file
        f = open(filename);

        r_table = [];
        V_table = [];
        F_table = [];

        # read in lines from the file
        for line in f.readlines():
            line = line.strip();

            # skip comment lines
            if line[0] == '#':
                continue;

            # split out the columns
            cols = line.split();
            values = [float(f) for f in cols];

            # validate the input
            if len(values) != 3:
                hoomd.context.current.device.cpp_msg.error("bond.table: file must have exactly 3 columns\n");
                raise RuntimeError("Error reading table file");

            # append to the tables
            r_table.append(values[0]);
            V_table.append(values[1]);
            F_table.append(values[2]);

        # validate input
        if self.width != len(r_table):
            hoomd.context.current.device.cpp_msg.error("bond.table: file must have exactly " + str(self.width) + " rows\n");
            raise RuntimeError("Error reading table file");

        # extract rmin and rmax
        rmin_table = r_table[0];
        rmax_table = r_table[-1];

        # check for even spacing
        dr = (rmax_table - rmin_table) / float(self.width-1);
        for i in range(0,self.width):
            r = rmin_table + dr * i;
            if math.fabs(r - r_table[i]) > 1e-3:
                hoomd.context.current.device.cpp_msg.error("bond.table: r must be monotonically increasing and evenly spaced\n");
                raise RuntimeError("Error reading table file");

        self.bond_coeff.set(bondname, func=_table_eval, rmin=rmin_table, rmax=rmax_table, coeff=dict(V=V_table, F=F_table, width=self.width))
<|endoftext|>


<commit_before><commit_msg>رفع باگ<commit_after><|endoftext|>


<commit_before># encoding: utf-8

#   Copyright 2012 Red Hat, Inc.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import logging
import sys
import os
import os.path
import json
from Singleton import Singleton

PLUGIN_TYPES = ('OS', 'CLOUD')
INFO_FILE_EXTENSION = '.info'
PKG_STR = 'imagefactory-plugins'

class PluginManager(Singleton):
    """ Registers and manages ImageFactory plugins. """
    @property
    def plugins(self):
        """
        The property plugins
        """
        return self._plugins

    def _singleton_init(self, plugin_path):
        self.log = logging.getLogger('%s.%s' % (__name__, self.__class__.__name__))
        sys.path.append(plugin_path)

        if(os.path.exists(plugin_path)):
            self.path = plugin_path
        else:
            msg = 'Plugin path (%s) does not exist! No plugins loaded.' % plugin_path
            self.log.exception(msg)
            raise Exception(msg)

        self._plugins = dict()
        self._targets = dict()
        self._types = dict().fromkeys(PLUGIN_TYPES, list())

    def load(self):
        """
        Enumerates through installed plugins and registers each according to
        the features provided. Only one plugin may be registered per feature.
        When more than one plugin is found, the first will be registered and
        all others listed as inactive.
        """

        info_files = list()
        directory_listing = os.listdir(self.path)
        for _file in directory_listing:
            if _file.endswith(INFO_FILE_EXTENSION):
                info_files.append(_file)

        for filename in info_files:
            plugin_name = filename.rstrip(INFO_FILE_EXTENSION)
            md = self.metadata_for_plugin(plugin_name)
            try:
                if(md['type'].upper() in PLUGIN_TYPES):
                    for target in md['targets']:
                        target = target if isinstance(target, str) else tuple(target)
                        if(not target in self._targets):
                            self._targets[target] = plugin_name
                        else:
                            msg = 'Did not register %s for %s. Plugin %s already registered.' % (plugin_name, target, self._targets[target])
                            self._register_plugin_with_error(plugin_name, msg)
                            self.log.warn(msg)
                    self._plugins[plugin_name] = md
                    self._types[md['type'].upper()].append(plugin_name)
                    self.log.info('Plugin (%s) loaded...' % plugin_name)
            except KeyError as e:
                msg = 'Invalid metadata for plugin (%s). Missing entry for %s.' % (plugin_name, e)
                self._register_plugin_with_error(plugin_name, msg)
                self.log.exception(msg)
            except Exception as e:
                msg = 'Loading plugin (%s) failed with exception: %s' % (plugin_name, e)
                self._register_plugin_with_error(plugin_name, msg)
                self.log.exception(msg)

    def _register_plugin_with_error(self, plugin_name, error_msg):
        self._plugins[plugin_name] = dict(ERROR = error_msg)

    def metadata_for_plugin(self, plugin):
        """
        Returns the metadata dictionary for the plugin.

        @param plugin name of the plugin or the plugin's info file

        @return dictionary containing the plugin's metadata
        """
        if(plugin in self._plugins):
            return self._plugins[plugin]
        else:
            fp = None
            metadata = None
            info_file = plugin + INFO_FILE_EXTENSION
            try:
                fp = open(os.path.join(self.path, info_file), 'r')
                metadata = json.load(fp)
            except Exception as e:
                self.log.exception('Exception caught while loading plugin metadata: %s' % e)
                raise e
            finally:
                if (fp):
                    fp.close()
                return metadata

    def plugin_for_target(self, target):
        """
        Looks up the plugin for a given target and returns an instance of the
        delegate class or None if no plugin is registered for the given target.
        Matches are done from left to right, ie. ('Fedora', '16', 'x86_64') will
        match a plugin with a target of ('Fedora', None, None) but not
        ('Fedora', None, 'x86_64')

        @param target A list or string matching the target field of the
        plugin's .info file.

        @return An instance of the delegate class of the plugin or None.
        """
        plugin_name = None
        try:
            if isinstance(target, str):
                self.log.debug("Attempting to match string target (%s)" % (target))
                plugin_name = self._targets.get(tuple([ target ]))
                plugin = __import__('%s.%s' % (PKG_STR, plugin_name), fromlist=['delegate_class'])
                return plugin.delegate_class()
            elif(isinstance(target, tuple)):
                _target = list(target)
                self.log.debug("Attempting to match list target (%s)" % (str(_target)))
                for index in range(len(target)):
                    plugin_name = self._targets.get(tuple(_target))
                    if(not plugin_name):
                        _target[-index] = None
                    else:
                        plugin = __import__('%s.%s' % (PKG_STR, plugin_name), fromlist=['delegate_class'])
                        return plugin.delegate_class()
        except Exception as e:
                self.log.exception('Exception caught during plugin lookup: %s' % e)
                return None
<commit_msg>The original range statement produces a range that starts with zero.<commit_after># encoding: utf-8

#   Copyright 2012 Red Hat, Inc.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import logging
import sys
import os
import os.path
import json
from Singleton import Singleton

PLUGIN_TYPES = ('OS', 'CLOUD')
INFO_FILE_EXTENSION = '.info'
PKG_STR = 'imagefactory-plugins'

class PluginManager(Singleton):
    """ Registers and manages ImageFactory plugins. """
    @property
    def plugins(self):
        """
        The property plugins
        """
        return self._plugins

    def _singleton_init(self, plugin_path):
        self.log = logging.getLogger('%s.%s' % (__name__, self.__class__.__name__))
        sys.path.append(plugin_path)

        if(os.path.exists(plugin_path)):
            self.path = plugin_path
        else:
            msg = 'Plugin path (%s) does not exist! No plugins loaded.' % plugin_path
            self.log.exception(msg)
            raise Exception(msg)

        self._plugins = dict()
        self._targets = dict()
        self._types = dict().fromkeys(PLUGIN_TYPES, list())

    def load(self):
        """
        Enumerates through installed plugins and registers each according to
        the features provided. Only one plugin may be registered per feature.
        When more than one plugin is found, the first will be registered and
        all others listed as inactive.
        """

        info_files = list()
        directory_listing = os.listdir(self.path)
        for _file in directory_listing:
            if _file.endswith(INFO_FILE_EXTENSION):
                info_files.append(_file)

        for filename in info_files:
            plugin_name = filename.rstrip(INFO_FILE_EXTENSION)
            md = self.metadata_for_plugin(plugin_name)
            try:
                if(md['type'].upper() in PLUGIN_TYPES):
                    for target in md['targets']:
                        target = target if isinstance(target, str) else tuple(target)
                        if(not target in self._targets):
                            self._targets[target] = plugin_name
                        else:
                            msg = 'Did not register %s for %s. Plugin %s already registered.' % (plugin_name, target, self._targets[target])
                            self._register_plugin_with_error(plugin_name, msg)
                            self.log.warn(msg)
                    self._plugins[plugin_name] = md
                    self._types[md['type'].upper()].append(plugin_name)
                    self.log.info('Plugin (%s) loaded...' % plugin_name)
            except KeyError as e:
                msg = 'Invalid metadata for plugin (%s). Missing entry for %s.' % (plugin_name, e)
                self._register_plugin_with_error(plugin_name, msg)
                self.log.exception(msg)
            except Exception as e:
                msg = 'Loading plugin (%s) failed with exception: %s' % (plugin_name, e)
                self._register_plugin_with_error(plugin_name, msg)
                self.log.exception(msg)

    def _register_plugin_with_error(self, plugin_name, error_msg):
        self._plugins[plugin_name] = dict(ERROR = error_msg)

    def metadata_for_plugin(self, plugin):
        """
        Returns the metadata dictionary for the plugin.

        @param plugin name of the plugin or the plugin's info file

        @return dictionary containing the plugin's metadata
        """
        if(plugin in self._plugins):
            return self._plugins[plugin]
        else:
            fp = None
            metadata = None
            info_file = plugin + INFO_FILE_EXTENSION
            try:
                fp = open(os.path.join(self.path, info_file), 'r')
                metadata = json.load(fp)
            except Exception as e:
                self.log.exception('Exception caught while loading plugin metadata: %s' % e)
                raise e
            finally:
                if (fp):
                    fp.close()
                return metadata

    def plugin_for_target(self, target):
        """
        Looks up the plugin for a given target and returns an instance of the
        delegate class or None if no plugin is registered for the given target.
        Matches are done from left to right, ie. ('Fedora', '16', 'x86_64') will
        match a plugin with a target of ('Fedora', None, None) but not
        ('Fedora', None, 'x86_64')

        @param target A list or string matching the target field of the
        plugin's .info file.

        @return An instance of the delegate class of the plugin or None.
        """
        plugin_name = None
        try:
            if isinstance(target, str):
                self.log.debug("Attempting to match string target (%s)" % (target))
                plugin_name = self._targets.get(tuple([ target ]))
                plugin = __import__('%s.%s' % (PKG_STR, plugin_name), fromlist=['delegate_class'])
                return plugin.delegate_class()
            elif(isinstance(target, tuple)):
                _target = list(target)
                self.log.debug("Attempting to match list target (%s)" % (str(_target)))
                for index in range(1,len(target)+1):
                    plugin_name = self._targets.get(tuple(_target))
                    if(not plugin_name):
                        _target[-index] = None
                    else:
                        plugin = __import__('%s.%s' % (PKG_STR, plugin_name), fromlist=['delegate_class'])
                        return plugin.delegate_class()
        except Exception as e:
                self.log.exception('Exception caught during plugin lookup: %s' % e)
                return None
<|endoftext|>


<commit_before>"""Constructs deep network reward models."""

from abc import ABC, abstractmethod

import tensorflow as tf

import imitation.util as util

class RewardNet(ABC):
  """Abstract reward network.

  This class assumes that the caller will set the default TensorFlow Session
  and initialize the network's variables.

  Attributes:
    env (gym.Env): the environment we are predicting reward for.
    old_obs_ph (tf.Tensor): previous observation placeholder.
    act_ph (tf.Tensor): action placeholder.
    new_obs_ph (tf.Tensor): next observation placeholder.
  """

  def __init__(self, env):
    """Builds a reward network.

    Args:
        env (gym.Env or str): The environment we are predicting reward for.
    """

    self.env = util.maybe_load_env(env)

    phs = util.build_placeholders(self.env, True)
    self.old_obs_ph, self.act_ph, self.new_obs_ph = phs

    with tf.variable_scope("theta_network"):
      self._theta_output = self.build_theta_network(
          self.old_obs_ph, self.act_ph)
      # TODO: assert that all the outputs above have shape (None,).

  @property
  @abstractmethod
  def reward_output_train(self):
    """A Tensor holding the training reward associated with each timestep.

    Different concrete subclasses will require different placeholders to be
    filled to calculate this output, but for all subclasses, filling the
    following placeholders will be sufficient:

    ```
    self.old_obs_ph
    self.act_ph
    self.new_obs_ph
    ```

    Returns:
      _reward_output_train (tf.Tensor) A (None,): shaped Tensor holding
          the training reward associated with each timestep.
    """
    pass

  @property
  def reward_output_test(self):
    """A Tensor holding the test reward associated with each timestep.

    Note this is the reward we use for transfer learning.

    Different concrete subclasses will require different
    placeholders to be filled to calculate this output, but for all
    subclasses, filling the following placeholders will be sufficient:

    ```
    self.old_obs_ph
    self.act_ph
    self.new_obs_ph
    ```

    Returns:
      _reward_output_test (tf.Tensor) A (None,): shaped Tensor holding
        the test reward associated with each timestep.
    """
    return self._theta_output

  @abstractmethod
  def build_theta_network(self, obs_input, act_input):
    """Builds the test reward network.

    The output of the network is the same as the reward used for transfer
    learning, and is the Tensor returned by `self.reward_output_test()`.

    Args:
      obs_input (tf.Tensor): The observation input. Its shape is
          `((None,) + self.env.observation_space.shape)`.
      act_input (tf.Tensor): The action input. Its shape is
          `((None,) + self.env.action_space.shape)`. The None dimension is
          expected to be the same as None dimension from `obs_input`.

    Returns:
      theta_output (tf.Tensor): A reward prediction for each of the
          inputs. The shape is `(None,)`.
    """
    pass

  def build_summaries(self):
    tf.summary.histogram("train_reward", self.reward_output_train)
    tf.summary.histogram("test_reward", self.reward_output_test)

class RewardNetShaped(RewardNet):
  """Abstract reward network with a phi network to shape training reward.

  This RewardNet formulation matches Equation (4) in the AIRL paper.
  Note that the experiments in Table 2 of the same paper showed shaped
  training rewards to be inferior to an unshaped training rewards in
  a Pendulum environment imitation learning task (and maybe HalfCheetah).
  (See original implementation of Pendulum experiment's reward function at
  https://github.com/justinjfu/inverse_rl/blob/master/inverse_rl/models/imitation_learning.py#L374)

  To make a concrete subclass, implement `build_phi_network()` and
  `build_theta_network()`.
  """

  def __init__(self, env, *, discount_factor=0.9, **kwargs):
    super().__init__(env, **kwargs)
    self._discount_factor = discount_factor

    with tf.variable_scope("phi_network"):
      res = self.build_phi_network(self.old_obs_ph, self.new_obs_ph)
      self._old_shaping_output, self._new_shaping_output = res

    with tf.variable_scope("f_network"):
      self._shaped_reward_output = (
        self._theta_output
        + self._discount_factor * self._new_shaping_output
        - self._old_shaping_output)

  @property
  def reward_output_train(self):
    """A Tensor holding the (shaped) training reward of each timestep.

    Requires the following placeholders to be filled:

    ```
    self.old_obs_ph
    self.act_ph
    self.new_obs_ph
    ```

    Returns:
      _reward_output_train (tf.Tensor) A (None,): shaped Tensor holding
          the training reward associated with each timestep.
    """
    return self._shaped_reward_output

  @abstractmethod
  def build_phi_network(self, old_obs_input, new_obs_input):
    """Build the reward shaping network (disentangles dynamics from reward).

    XXX: We could potentially make it easier on the subclasser by requiring
    only one input. ie build_phi_network(obs_input). Later in
    _build_f_network, I could stack Tensors of old and new observations,
    pass them simulatenously through the network, and then unstack the
    outputs. Another way to do this would be to pass in a single
    rank 3 obs_input with shape `(2, None) + self.env.observation_space`.

    Args:
      old_obs_input (tf.Tensor): The old observations (corresponding to the
          state at which the current action is made). The shape of this
          Tensor should be `(None,) + self.env.observation_space.shape`.
      new_obs_input (tf.Tensor): The new observations (corresponding to
          the state that we transition to after this state-action pair.

    Returns:
      old_shaping_output (tf.Tensor): A reward shaping prediction for
          each of the old observation inputs. Has shape `(None,)`.
      new_shaping_output (tf.Tensor): A reward shaping prediction for
          each of the new observation inputs. Has shape `(None,)`
    """
    pass

  def build_summaries(self):
    super().build_summaries()
    tf.summary.histogram("shaping_old", self._old_shaping_output)
    tf.summary.histogram("shaping_new", self._new_shaping_output)

def build_basic_theta_network(hid_sizes, obs_space, act_space,
                              old_obs_input, new_obs_input, act_input):
  """Builds a reward network depending on specified observations and actions.

  All specified inputs will be preprocessed and then concatenated. If all
  inputs are specified, then it will be a :math:`R(o,a,o')` network.
  Conversely, if `new_obs_input` and `act_input` are both set to `None`, it
  will depend just on the current observation: :math:`R(o)`.

  Arguments:
    hid_sizes (Optional[List[int]]): Number of units at each hidden layer.
    obs_space (gym.Space): Observation space.
    act_space (gym.Space): Action space.
    old_obs_input (Optional[tf.Tensor]): Previous observation.
    new_obs_input (Optional[tf.Tensor]): Next observation.
    act_input (tf.Tensor): Action.

  Returns:
    tf.Tensor: Predicted reward.

  Raises:
    ValueError: If all of old_obs_input, new_obs_input and act_input are None.
  """
  if hid_sizes is None:
    hid_sizes = []

  with tf.variable_scope("theta"):
    inputs = [(old_obs_input, obs_space),
              (act_input, act_space),
              (new_obs_input, obs_space)]
    inputs = [(input, space) for input, space in inputs if input is not None]
    if len(inputs) == 0:
      raise ValueError("Must specify at least one input")

    inputs = [util.flat(input, space.shape) for input, space in inputs]
    inputs = tf.concat(inputs, axis=1)
    theta_output = util.apply_ff(inputs, hid_sizes=hid_sizes)

    return theta_output

class BasicRewardNet(RewardNet):
  """An unshaped reward net with simple, default settings.

  Intended to match the reward network trained for
  the original AIRL pendulum experiments. Right now it has a linear function
  approximator for the theta network, not sure if this is what I want.
  """

  def __init__(self, env, *, state_only=False, theta_units=None, **kwargs):
    """Builds a simple reward network.

    Args:
      env (gym.Env or str): The environment we are predicting reward for.
      state_only (bool): If True, then ignore the action when predicting
          and training the reward network phi.
      theta_units (List[int]): number of hidden units at each layer of
          the feedforward theta network.
    """
    self.state_only = state_only
    self.theta_units = theta_units
    super().__init__(env, **kwargs)

  def build_theta_network(self, obs_input, act_input):
    act_or_none = None if self.state_only else act_input
    return build_basic_theta_network(self.theta_units,
                                     self.env.observation_space,
                                     self.env.action_space,
                                     old_obs_input=obs_input,
                                     act_input=act_or_none,
                                     new_obs_input=None)

  @property
  def reward_output_train(self):
    """Training reward is the same as the test reward, since no shaping."""
    return self.reward_output_test

def build_basic_phi_network(hid_sizes, obs_space, old_obs_input, new_obs_input):
  """Builds a potential network depending on specified observation.

  Arguments:
    hid_sizes (Optional[List[int]]): Number of units at each hidden layer.
    obs_space (gym.Space): Observation space.
    old_obs_input (Optional[tf.Tensor]): Previous observation.
    new_obs_input (Optional[tf.Tensor]): Next observation.

  Returns:
    Tuple[tf.Tensor]: potential for the old and new observations.
  """
  if hid_sizes is None:
    hid_sizes = [32, 32]

  with tf.variable_scope("phi", reuse=tf.AUTO_REUSE):
    old_o = util.flat(old_obs_input, obs_space)
    new_o = util.flat(new_obs_input, obs_space)

    # Weight share, just with different inputs old_o and new_o
    old_shaping_output = tf.identity(
        util.apply_ff(old_o, hid_sizes=hid_sizes),
        name="old_shaping_output")
    new_shaping_output = tf.identity(
        util.apply_ff(new_o, hid_sizes=hid_sizes),
        name="new_shaping_output")

  return old_shaping_output, new_shaping_output

class BasicShapedRewardNet(RewardNetShaped):
  """A shaped reward network with simple, default settings.

  With default parameters this RewardNet has two hidden layers [32, 32]
  for the reward shaping phi network, and a linear function approximator
  for the theta network. These settings match the network architectures for
  continuous control experiments described in Appendix D.1 of the
  AIRL paper.

  This network flattens inputs. So it isn't suitable for training on
  pixel observations.
  """

  def __init__(self, env, *, state_only=False,
               theta_units=None, phi_units=None, **kwargs):
    """Builds a simple shaped reward network.

    Args:
      env (gym.Env or str): The environment we are predicitng reward for.
      units (int): The number of hidden units in each of the two layers
          of the feed forward phi network.
      state_only (bool): If True, then ignore the action when predicting
          and training the reward network phi.
      theta_units (List[int]): number of hidden units at each layer of
          the feedforward theta network.
      phi_units: (List[int] number of hidden units at each layer of
          the feedforward phi network.
      discount_factor (float): A number in the range [0, 1].
    """
    self.state_only = state_only
    self.theta_units = theta_units
    self.phi_units = phi_units
    super().__init__(env, **kwargs)

  def build_theta_network(self, obs_input, act_input):
    act_or_none = None if self.state_only else act_input
    return build_basic_theta_network(self.theta_units,
                                     self.env.observation_space,
                                     self.env.action_space,
                                     old_obs_input=obs_input,
                                     act_input=act_or_none,
                                     new_obs_input=None)

  def build_phi_network(self, old_obs_input, new_obs_input):
    return build_basic_phi_network(self.phi_units, self.env.observation_space,
                                   old_obs_input, new_obs_input)
<commit_msg>Bugfix: use shapes in util.flat<commit_after>"""Constructs deep network reward models."""

from abc import ABC, abstractmethod

import tensorflow as tf

import imitation.util as util

class RewardNet(ABC):
  """Abstract reward network.

  This class assumes that the caller will set the default TensorFlow Session
  and initialize the network's variables.

  Attributes:
    env (gym.Env): the environment we are predicting reward for.
    old_obs_ph (tf.Tensor): previous observation placeholder.
    act_ph (tf.Tensor): action placeholder.
    new_obs_ph (tf.Tensor): next observation placeholder.
  """

  def __init__(self, env):
    """Builds a reward network.

    Args:
        env (gym.Env or str): The environment we are predicting reward for.
    """

    self.env = util.maybe_load_env(env)

    phs = util.build_placeholders(self.env, True)
    self.old_obs_ph, self.act_ph, self.new_obs_ph = phs

    with tf.variable_scope("theta_network"):
      self._theta_output = self.build_theta_network(
          self.old_obs_ph, self.act_ph)
      # TODO: assert that all the outputs above have shape (None,).

  @property
  @abstractmethod
  def reward_output_train(self):
    """A Tensor holding the training reward associated with each timestep.

    Different concrete subclasses will require different placeholders to be
    filled to calculate this output, but for all subclasses, filling the
    following placeholders will be sufficient:

    ```
    self.old_obs_ph
    self.act_ph
    self.new_obs_ph
    ```

    Returns:
      _reward_output_train (tf.Tensor) A (None,): shaped Tensor holding
          the training reward associated with each timestep.
    """
    pass

  @property
  def reward_output_test(self):
    """A Tensor holding the test reward associated with each timestep.

    Note this is the reward we use for transfer learning.

    Different concrete subclasses will require different
    placeholders to be filled to calculate this output, but for all
    subclasses, filling the following placeholders will be sufficient:

    ```
    self.old_obs_ph
    self.act_ph
    self.new_obs_ph
    ```

    Returns:
      _reward_output_test (tf.Tensor) A (None,): shaped Tensor holding
        the test reward associated with each timestep.
    """
    return self._theta_output

  @abstractmethod
  def build_theta_network(self, obs_input, act_input):
    """Builds the test reward network.

    The output of the network is the same as the reward used for transfer
    learning, and is the Tensor returned by `self.reward_output_test()`.

    Args:
      obs_input (tf.Tensor): The observation input. Its shape is
          `((None,) + self.env.observation_space.shape)`.
      act_input (tf.Tensor): The action input. Its shape is
          `((None,) + self.env.action_space.shape)`. The None dimension is
          expected to be the same as None dimension from `obs_input`.

    Returns:
      theta_output (tf.Tensor): A reward prediction for each of the
          inputs. The shape is `(None,)`.
    """
    pass

  def build_summaries(self):
    tf.summary.histogram("train_reward", self.reward_output_train)
    tf.summary.histogram("test_reward", self.reward_output_test)

class RewardNetShaped(RewardNet):
  """Abstract reward network with a phi network to shape training reward.

  This RewardNet formulation matches Equation (4) in the AIRL paper.
  Note that the experiments in Table 2 of the same paper showed shaped
  training rewards to be inferior to an unshaped training rewards in
  a Pendulum environment imitation learning task (and maybe HalfCheetah).
  (See original implementation of Pendulum experiment's reward function at
  https://github.com/justinjfu/inverse_rl/blob/master/inverse_rl/models/imitation_learning.py#L374)

  To make a concrete subclass, implement `build_phi_network()` and
  `build_theta_network()`.
  """

  def __init__(self, env, *, discount_factor=0.9, **kwargs):
    super().__init__(env, **kwargs)
    self._discount_factor = discount_factor

    with tf.variable_scope("phi_network"):
      res = self.build_phi_network(self.old_obs_ph, self.new_obs_ph)
      self._old_shaping_output, self._new_shaping_output = res

    with tf.variable_scope("f_network"):
      self._shaped_reward_output = (
        self._theta_output
        + self._discount_factor * self._new_shaping_output
        - self._old_shaping_output)

  @property
  def reward_output_train(self):
    """A Tensor holding the (shaped) training reward of each timestep.

    Requires the following placeholders to be filled:

    ```
    self.old_obs_ph
    self.act_ph
    self.new_obs_ph
    ```

    Returns:
      _reward_output_train (tf.Tensor) A (None,): shaped Tensor holding
          the training reward associated with each timestep.
    """
    return self._shaped_reward_output

  @abstractmethod
  def build_phi_network(self, old_obs_input, new_obs_input):
    """Build the reward shaping network (disentangles dynamics from reward).

    XXX: We could potentially make it easier on the subclasser by requiring
    only one input. ie build_phi_network(obs_input). Later in
    _build_f_network, I could stack Tensors of old and new observations,
    pass them simulatenously through the network, and then unstack the
    outputs. Another way to do this would be to pass in a single
    rank 3 obs_input with shape `(2, None) + self.env.observation_space`.

    Args:
      old_obs_input (tf.Tensor): The old observations (corresponding to the
          state at which the current action is made). The shape of this
          Tensor should be `(None,) + self.env.observation_space.shape`.
      new_obs_input (tf.Tensor): The new observations (corresponding to
          the state that we transition to after this state-action pair.

    Returns:
      old_shaping_output (tf.Tensor): A reward shaping prediction for
          each of the old observation inputs. Has shape `(None,)`.
      new_shaping_output (tf.Tensor): A reward shaping prediction for
          each of the new observation inputs. Has shape `(None,)`
    """
    pass

  def build_summaries(self):
    super().build_summaries()
    tf.summary.histogram("shaping_old", self._old_shaping_output)
    tf.summary.histogram("shaping_new", self._new_shaping_output)

def build_basic_theta_network(hid_sizes, obs_space, act_space,
                              old_obs_input, new_obs_input, act_input):
  """Builds a reward network depending on specified observations and actions.

  All specified inputs will be preprocessed and then concatenated. If all
  inputs are specified, then it will be a :math:`R(o,a,o')` network.
  Conversely, if `new_obs_input` and `act_input` are both set to `None`, it
  will depend just on the current observation: :math:`R(o)`.

  Arguments:
    hid_sizes (Optional[List[int]]): Number of units at each hidden layer.
    obs_space (gym.Space): Observation space.
    act_space (gym.Space): Action space.
    old_obs_input (Optional[tf.Tensor]): Previous observation.
    new_obs_input (Optional[tf.Tensor]): Next observation.
    act_input (tf.Tensor): Action.

  Returns:
    tf.Tensor: Predicted reward.

  Raises:
    ValueError: If all of old_obs_input, new_obs_input and act_input are None.
  """
  if hid_sizes is None:
    hid_sizes = []

  with tf.variable_scope("theta"):
    inputs = [(old_obs_input, obs_space),
              (act_input, act_space),
              (new_obs_input, obs_space)]
    inputs = [(input, space) for input, space in inputs if input is not None]
    if len(inputs) == 0:
      raise ValueError("Must specify at least one input")

    inputs = [util.flat(input, space.shape) for input, space in inputs]
    inputs = tf.concat(inputs, axis=1)
    theta_output = util.apply_ff(inputs, hid_sizes=hid_sizes)

    return theta_output

class BasicRewardNet(RewardNet):
  """An unshaped reward net with simple, default settings.

  Intended to match the reward network trained for
  the original AIRL pendulum experiments. Right now it has a linear function
  approximator for the theta network, not sure if this is what I want.
  """

  def __init__(self, env, *, state_only=False, theta_units=None, **kwargs):
    """Builds a simple reward network.

    Args:
      env (gym.Env or str): The environment we are predicting reward for.
      state_only (bool): If True, then ignore the action when predicting
          and training the reward network phi.
      theta_units (List[int]): number of hidden units at each layer of
          the feedforward theta network.
    """
    self.state_only = state_only
    self.theta_units = theta_units
    super().__init__(env, **kwargs)

  def build_theta_network(self, obs_input, act_input):
    act_or_none = None if self.state_only else act_input
    return build_basic_theta_network(self.theta_units,
                                     self.env.observation_space,
                                     self.env.action_space,
                                     old_obs_input=obs_input,
                                     act_input=act_or_none,
                                     new_obs_input=None)

  @property
  def reward_output_train(self):
    """Training reward is the same as the test reward, since no shaping."""
    return self.reward_output_test

def build_basic_phi_network(hid_sizes, obs_space, old_obs_input, new_obs_input):
  """Builds a potential network depending on specified observation.

  Arguments:
    hid_sizes (Optional[List[int]]): Number of units at each hidden layer.
    obs_space (gym.Space): Observation space.
    old_obs_input (Optional[tf.Tensor]): Previous observation.
    new_obs_input (Optional[tf.Tensor]): Next observation.

  Returns:
    Tuple[tf.Tensor]: potential for the old and new observations.
  """
  if hid_sizes is None:
    hid_sizes = [32, 32]

  with tf.variable_scope("phi", reuse=tf.AUTO_REUSE):
    old_o = util.flat(old_obs_input, obs_space.shape)
    new_o = util.flat(new_obs_input, obs_space.shape)

    # Weight share, just with different inputs old_o and new_o
    old_shaping_output = tf.identity(
        util.apply_ff(old_o, hid_sizes=hid_sizes),
        name="old_shaping_output")
    new_shaping_output = tf.identity(
        util.apply_ff(new_o, hid_sizes=hid_sizes),
        name="new_shaping_output")

  return old_shaping_output, new_shaping_output

class BasicShapedRewardNet(RewardNetShaped):
  """A shaped reward network with simple, default settings.

  With default parameters this RewardNet has two hidden layers [32, 32]
  for the reward shaping phi network, and a linear function approximator
  for the theta network. These settings match the network architectures for
  continuous control experiments described in Appendix D.1 of the
  AIRL paper.

  This network flattens inputs. So it isn't suitable for training on
  pixel observations.
  """

  def __init__(self, env, *, state_only=False,
               theta_units=None, phi_units=None, **kwargs):
    """Builds a simple shaped reward network.

    Args:
      env (gym.Env or str): The environment we are predicitng reward for.
      units (int): The number of hidden units in each of the two layers
          of the feed forward phi network.
      state_only (bool): If True, then ignore the action when predicting
          and training the reward network phi.
      theta_units (List[int]): number of hidden units at each layer of
          the feedforward theta network.
      phi_units: (List[int] number of hidden units at each layer of
          the feedforward phi network.
      discount_factor (float): A number in the range [0, 1].
    """
    self.state_only = state_only
    self.theta_units = theta_units
    self.phi_units = phi_units
    super().__init__(env, **kwargs)

  def build_theta_network(self, obs_input, act_input):
    act_or_none = None if self.state_only else act_input
    return build_basic_theta_network(self.theta_units,
                                     self.env.observation_space,
                                     self.env.action_space,
                                     old_obs_input=obs_input,
                                     act_input=act_or_none,
                                     new_obs_input=None)

  def build_phi_network(self, old_obs_input, new_obs_input):
    return build_basic_phi_network(self.phi_units, self.env.observation_space,
                                   old_obs_input, new_obs_input)
<|endoftext|>