<commit_before>import logging
import multiprocessing
import os
import shutil

from fusesoc.config import Config
from fusesoc import utils
from fusesoc.utils import pr_info
from fusesoc.core import OptionSectionMissing
from .simulator import Simulator

logger = logging.getLogger(__name__)

class Source(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class Verilator(Simulator):

    def __init__(self, system):
        super(Verilator, self).__init__(system)

        self.verilator_options = []
        self.src_files = []
        self.include_files = []
        self.include_dirs = []
        self.tb_toplevel = ""
        self.src_type = ''
        self.define_files = []
        self.libs = []
        self.top_module = ""

        if system.verilator is not None:
             self.verilator_options  = system.verilator.verilator_options
             self.src_files          = system.verilator.src_files
             self.include_files      = system.verilator.include_files
             self.tb_toplevel        = system.verilator.tb_toplevel
             self.src_type           = system.verilator.source_type
             self.define_files       = system.verilator.define_files
             self.top_module         = system.verilator.top_module
             self.fusesoc_cli_parser = system.verilator.cli_parser and (system.verilator.cli_parser == 'fusesoc')

        if self.top_module == '':
            raise OptionSectionMissing('top_module')
        if self.tb_toplevel == '':
            raise OptionSectionMissing('tb_toplevel')

    def export(self):
        src_dir = self.system.files_root
        dst_dir = os.path.join(self.src_root, self.system.sanitized_name)
        src_files = [f.name for f in self.src_files]
        src_files += [f.name for f in self.include_files]
        src_files += [self.tb_toplevel]
        dirs = list(set(map(os.path.dirname, src_files)))
        for d in dirs:
            if not os.path.exists(os.path.join(dst_dir, d)):
                os.makedirs(os.path.join(dst_dir, d))

        for f in src_files:
            if(os.path.exists(os.path.join(src_dir, f))):
                shutil.copyfile(os.path.join(src_dir, f),
                                os.path.join(dst_dir, f))

    def configure(self, args):
        skip = not self.fusesoc_cli_parser
        super(Verilator, self).configure(args, skip_params = skip)
        self.export()
        self._write_config_files()
        #self.object_files = [os.path.splitext(os.path.basename(s))[0]+'.o' for s in self.src_files]

    def _write_config_files(self):
        self.verilator_file = 'input.vc'
        f = open(os.path.join(self.work_root,self.verilator_file),'w')

        incdirs = set()
        src_files = []

        (src_files, incdirs) = self._get_fileset_files(['synth', 'verilator'])
        for include_dir in incdirs:
            f.write("+incdir+" + include_dir + '\n')
        for src_file in src_files:
            f.write(src_file.name + '\n')
        f.close()
        #convert verilog defines into C file
        for files in self.define_files:
            read_file = os.path.join(self.src_root,files)
            write_file = os.path.join(os.path.dirname(os.path.join(self.work_root,self.tb_toplevel)),os.path.splitext(os.path.basename(files))[0]+'.h')
            utils.convert_V2H(read_file, write_file)

    def _verilate(self):
        if self.src_type == 'systemC':
            args = ['--sc']
        else:
            args = ['--cc']
        args += ['-f', self.verilator_file]
        args += ['--top-module', self.top_module]
        args += ['--exe']
        args += ['-LDFLAGS "']
        args += ['-Wl,--start-group']
        args += [os.path.join(self.work_root, s) for s in self.archives]
        args += ['-Wl,--end-group']
        args += [l for l in self.libs]
        args += ['"']
        args += ['-CFLAGS ' + '-I' + i for i in self.include_dirs]
        args += [os.path.join(self.src_root, self.system.sanitized_name, self.tb_toplevel)]
        for key, value in self.vlogparam.items():
            args += ['-G{}={}'.format(key, value)]
        args += self.verilator_options

        cmd = utils.find_verilator()
        if cmd is None:
             raise RuntimeError("VERILATOR_ROOT not set and there is no verilator program in your PATH")
        cmd += ' ' + ' '.join(args)

        if (self.src_type == 'systemC') and not utils.check_systemc_env():
            raise RuntimeError("Need $SYSTEMC_LIBDIR and $SYSTEMC_INCLUDE in environment or when Verilator configured")

        l = utils.Launcher(cmd,
                           shell=True,
                           cwd = self.work_root,
                           stderr = open(os.path.join(self.work_root,'verilator.log'),'w'),
                           stdout = open(os.path.join(self.work_root, 'verilator.out.log'),'w')
        )
        print('')
        pr_info("Starting Verilator:")
        if Config().verbose:
             pr_info("  Verilator working dir: " + self.work_root)
             pr_info("  Verilator command: " + cmd)
        print('')
        l.run()

    def build(self):
        super(Verilator, self).build()
        self.archives = []
        for core in self.cores:
            if core.verilator:
                 if core.verilator.archive:
                      self.archives += [core.sanitized_name+'.a']
                 self.libs += core.verilator.libs
                 self.include_dirs += [os.path.join(self.src_root, core.sanitized_name, d) for d in core.verilator.include_dirs]
        self.include_dirs += [os.path.dirname(os.path.join(self.work_root, self.tb_toplevel))]

        self.include_dirs += [self.src_root]

        pr_info("Verilating source")
        self._verilate()
        for core in self.cores:
            if core.verilator:
                 self._build(core, self.work_root, self.src_root)

        # Do parallel builds with <number of cpus> * 2 jobs.
        make_job_count = multiprocessing.cpu_count() * 2

        pr_info("Building verilator executable:")
        args = ['-f', 'V' + self.top_module + '.mk', '-j', str(make_job_count), 'V' + self.top_module]
        l = utils.Launcher('make', args,
                       cwd=os.path.join(self.work_root, 'obj_dir'),
                       stdout = open(os.path.join(self.work_root,
                                                  'verilator.make.log'),'w'))
        if Config().verbose:
             pr_info("  Verilator executable working dir: "
                     + os.path.join(self.work_root, 'obj_dir'))
             pr_info("  Verilator executable command: make " + ' '.join(args))
        l.run()

    def _build(self, core, work_root, src_root):
        source_type = core.verilator.source_type
        if source_type == 'C' or source_type == '':
             self.build_C(core, work_root, src_root)
        elif source_type == 'CPP':
            self.build_CPP(core, work_root, src_root)
        elif source_type == 'systemC':
            self.build_SysC(core, work_root, src_root)
        else:
            raise Source(core.verilator.source_type)

        if core.verilator._object_files:
            args = []
            args += ['rvs']
            args += [core.sanitized_name+'.a']
            args += core.verilator._object_files
            l = utils.Launcher('ar', args,
                     cwd=work_root)
            if Config().verbose:
                pr_info("  linker working dir: " + work_root)
                pr_info("  linker command: ar " + ' '.join(args))
            l.run()
            print()

    def build_C(self, core, work_root, src_root):
        args = ['-c']
        args += ['-std=c99']
        args += ['-I'+src_root]
        args += ['-I'+os.path.join(src_root, core.sanitized_name, s) for s in core.verilator.include_dirs]
        for src_file in core.verilator.src_files:
            pr_info("Compiling " + src_file.name)
            l = utils.Launcher('gcc',
                     args + [os.path.join(src_root, core.sanitized_name, src_file.name)],
                         cwd=work_root,
                         stderr = open(os.path.join(work_root, 'gcc.err.log'),'a'),
                         stdout = open(os.path.join(work_root, 'gcc.out.log'),'a'))
            if Config().verbose:
                pr_info("  C compilation working dir: " + work_root)
                pr_info("  C compilation command: gcc " + ' '.join(args) + ' ' + os.path.join(src_root, core.sanitized_name, src_file.name))
            l.run()

    def build_CPP(self, core, work_root, src_root):
        verilator_root = utils.get_verilator_root()
        if verilator_root is None:
            verilator_root = utils.get_verilator_root()
        args = ['-c']
        args += ['-I'+src_root]
        args += ['-I'+os.path.join(src_root, core.sanitized_name, s) for s in self.include_dirs]
        args += ['-Iobj_dir']
        args += ['-I'+os.path.join(verilator_root,'include')]
        args += ['-I'+os.path.join(verilator_root,'include', 'vltstd')]
        for src_file in core.verilator.src_files:
            pr_info("Compiling " + src_file.name)
            l = utils.Launcher('g++', args + [os.path.join(src_root, core.sanitized_name, src_file.name)],
                         cwd=work_root,
                         stderr = open(os.path.join(work_root, 'g++.err.log'),'a'))
            if Config().verbose:
                pr_info("  C++ compilation working dir: " + work_root)
                pr_info("  C++ compilation command: g++ " + ' '.join(args) + ' ' + os.path.join(src_root, core.sanitized_name, src_file.name))
            l.run()

    def build_SysC(self, core, work_root, src_root):
        verilator_root = utils.get_verilator_root()
        args = ['-I.']
        args += ['-MMD']
        args += ['-I'+src_root]
        args += ['-I'+s for s in self.include_dirs]
        args += ['-Iobj_dir']
        args += ['-I'+os.path.join(verilator_root,'include')]
        args += ['-I'+os.path.join(verilator_root,'include', 'vltstd')]
        args += ['-DVL_PRINTF=printf']
        args += ['-DVM_TRACE=1']
        args += ['-DVM_COVERAGE=0']
        if os.getenv('SYSTEMC_INCLUDE'):
            args += ['-I'+os.getenv('SYSTEMC_INCLUDE')]
        if os.getenv('SYSTEMC'):
            args += ['-I'+os.path.join(os.getenv('SYSTEMC'),'include')]
        args += ['-Wno-deprecated']
        if os.getenv('SYSTEMC_CXX_FLAGS'):
             args += [os.getenv('SYSTEMC_CXX_FLAGS')]
        args += ['-c']
        args += ['-g']

        for src_file in core.verilator.src_files:
            pr_info("Compiling " + src_file.name)
            l = utils.Launcher('g++', args + [os.path.join(src_root, core.sanitized_name, src_file.name)],
                         cwd=work_root,
                         stderr = open(os.path.join(work_root, 'g++.err.log'),'a'))
            if Config().verbose:
                pr_info("  SystemC compilation working dir: " + work_root)
                pr_info("  SystemC compilation command: g++ " + ' '.join(args) + ' ' + os.path.join(src_root, core.sanitized_name, src_file.name))
            l.run()

    def run(self, args):
        if not self.fusesoc_cli_parser:
            self.plusarg = []
        super(Verilator, self).run(args)
        self.env = os.environ.copy()
        self.env['CORE_ROOT'] = os.path.abspath(self.system.core_root)
        self.env['BUILD_ROOT'] = os.path.abspath(self.build_root)
        self.env['SIM_ROOT'] = os.path.abspath(self.work_root)
        if self.fusesoc_cli_parser:
            _args = []
            for key, value in self.plusarg.items():
                _args += ['+{}={}'.format(key, value)]
        else:
            _args = args
        pr_info("Running simulation")
        utils.Launcher('./V' + self.top_module,
                       _args,
                       cwd=os.path.join(self.work_root, 'obj_dir'),
                       env = self.env).run()
<commit_msg>verilator.py: Remove unused env var BUILD_ROOT<commit_after>import logging
import multiprocessing
import os
import shutil

from fusesoc.config import Config
from fusesoc import utils
from fusesoc.utils import pr_info
from fusesoc.core import OptionSectionMissing
from .simulator import Simulator

logger = logging.getLogger(__name__)

class Source(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class Verilator(Simulator):

    def __init__(self, system):
        super(Verilator, self).__init__(system)

        self.verilator_options = []
        self.src_files = []
        self.include_files = []
        self.include_dirs = []
        self.tb_toplevel = ""
        self.src_type = ''
        self.define_files = []
        self.libs = []
        self.top_module = ""

        if system.verilator is not None:
             self.verilator_options  = system.verilator.verilator_options
             self.src_files          = system.verilator.src_files
             self.include_files      = system.verilator.include_files
             self.tb_toplevel        = system.verilator.tb_toplevel
             self.src_type           = system.verilator.source_type
             self.define_files       = system.verilator.define_files
             self.top_module         = system.verilator.top_module
             self.fusesoc_cli_parser = system.verilator.cli_parser and (system.verilator.cli_parser == 'fusesoc')

        if self.top_module == '':
            raise OptionSectionMissing('top_module')
        if self.tb_toplevel == '':
            raise OptionSectionMissing('tb_toplevel')

    def export(self):
        src_dir = self.system.files_root
        dst_dir = os.path.join(self.src_root, self.system.sanitized_name)
        src_files = [f.name for f in self.src_files]
        src_files += [f.name for f in self.include_files]
        src_files += [self.tb_toplevel]
        dirs = list(set(map(os.path.dirname, src_files)))
        for d in dirs:
            if not os.path.exists(os.path.join(dst_dir, d)):
                os.makedirs(os.path.join(dst_dir, d))

        for f in src_files:
            if(os.path.exists(os.path.join(src_dir, f))):
                shutil.copyfile(os.path.join(src_dir, f),
                                os.path.join(dst_dir, f))

    def configure(self, args):
        skip = not self.fusesoc_cli_parser
        super(Verilator, self).configure(args, skip_params = skip)
        self.export()
        self._write_config_files()
        #self.object_files = [os.path.splitext(os.path.basename(s))[0]+'.o' for s in self.src_files]

    def _write_config_files(self):
        self.verilator_file = 'input.vc'
        f = open(os.path.join(self.work_root,self.verilator_file),'w')

        incdirs = set()
        src_files = []

        (src_files, incdirs) = self._get_fileset_files(['synth', 'verilator'])
        for include_dir in incdirs:
            f.write("+incdir+" + include_dir + '\n')
        for src_file in src_files:
            f.write(src_file.name + '\n')
        f.close()
        #convert verilog defines into C file
        for files in self.define_files:
            read_file = os.path.join(self.src_root,files)
            write_file = os.path.join(os.path.dirname(os.path.join(self.work_root,self.tb_toplevel)),os.path.splitext(os.path.basename(files))[0]+'.h')
            utils.convert_V2H(read_file, write_file)

    def _verilate(self):
        if self.src_type == 'systemC':
            args = ['--sc']
        else:
            args = ['--cc']
        args += ['-f', self.verilator_file]
        args += ['--top-module', self.top_module]
        args += ['--exe']
        args += ['-LDFLAGS "']
        args += ['-Wl,--start-group']
        args += [os.path.join(self.work_root, s) for s in self.archives]
        args += ['-Wl,--end-group']
        args += [l for l in self.libs]
        args += ['"']
        args += ['-CFLAGS ' + '-I' + i for i in self.include_dirs]
        args += [os.path.join(self.src_root, self.system.sanitized_name, self.tb_toplevel)]
        for key, value in self.vlogparam.items():
            args += ['-G{}={}'.format(key, value)]
        args += self.verilator_options

        cmd = utils.find_verilator()
        if cmd is None:
             raise RuntimeError("VERILATOR_ROOT not set and there is no verilator program in your PATH")
        cmd += ' ' + ' '.join(args)

        if (self.src_type == 'systemC') and not utils.check_systemc_env():
            raise RuntimeError("Need $SYSTEMC_LIBDIR and $SYSTEMC_INCLUDE in environment or when Verilator configured")

        l = utils.Launcher(cmd,
                           shell=True,
                           cwd = self.work_root,
                           stderr = open(os.path.join(self.work_root,'verilator.log'),'w'),
                           stdout = open(os.path.join(self.work_root, 'verilator.out.log'),'w')
        )
        print('')
        pr_info("Starting Verilator:")
        if Config().verbose:
             pr_info("  Verilator working dir: " + self.work_root)
             pr_info("  Verilator command: " + cmd)
        print('')
        l.run()

    def build(self):
        super(Verilator, self).build()
        self.archives = []
        for core in self.cores:
            if core.verilator:
                 if core.verilator.archive:
                      self.archives += [core.sanitized_name+'.a']
                 self.libs += core.verilator.libs
                 self.include_dirs += [os.path.join(self.src_root, core.sanitized_name, d) for d in core.verilator.include_dirs]
        self.include_dirs += [os.path.dirname(os.path.join(self.work_root, self.tb_toplevel))]

        self.include_dirs += [self.src_root]

        pr_info("Verilating source")
        self._verilate()
        for core in self.cores:
            if core.verilator:
                 self._build(core, self.work_root, self.src_root)

        # Do parallel builds with <number of cpus> * 2 jobs.
        make_job_count = multiprocessing.cpu_count() * 2

        pr_info("Building verilator executable:")
        args = ['-f', 'V' + self.top_module + '.mk', '-j', str(make_job_count), 'V' + self.top_module]
        l = utils.Launcher('make', args,
                       cwd=os.path.join(self.work_root, 'obj_dir'),
                       stdout = open(os.path.join(self.work_root,
                                                  'verilator.make.log'),'w'))
        if Config().verbose:
             pr_info("  Verilator executable working dir: "
                     + os.path.join(self.work_root, 'obj_dir'))
             pr_info("  Verilator executable command: make " + ' '.join(args))
        l.run()

    def _build(self, core, work_root, src_root):
        source_type = core.verilator.source_type
        if source_type == 'C' or source_type == '':
             self.build_C(core, work_root, src_root)
        elif source_type == 'CPP':
            self.build_CPP(core, work_root, src_root)
        elif source_type == 'systemC':
            self.build_SysC(core, work_root, src_root)
        else:
            raise Source(core.verilator.source_type)

        if core.verilator._object_files:
            args = []
            args += ['rvs']
            args += [core.sanitized_name+'.a']
            args += core.verilator._object_files
            l = utils.Launcher('ar', args,
                     cwd=work_root)
            if Config().verbose:
                pr_info("  linker working dir: " + work_root)
                pr_info("  linker command: ar " + ' '.join(args))
            l.run()
            print()

    def build_C(self, core, work_root, src_root):
        args = ['-c']
        args += ['-std=c99']
        args += ['-I'+src_root]
        args += ['-I'+os.path.join(src_root, core.sanitized_name, s) for s in core.verilator.include_dirs]
        for src_file in core.verilator.src_files:
            pr_info("Compiling " + src_file.name)
            l = utils.Launcher('gcc',
                     args + [os.path.join(src_root, core.sanitized_name, src_file.name)],
                         cwd=work_root,
                         stderr = open(os.path.join(work_root, 'gcc.err.log'),'a'),
                         stdout = open(os.path.join(work_root, 'gcc.out.log'),'a'))
            if Config().verbose:
                pr_info("  C compilation working dir: " + work_root)
                pr_info("  C compilation command: gcc " + ' '.join(args) + ' ' + os.path.join(src_root, core.sanitized_name, src_file.name))
            l.run()

    def build_CPP(self, core, work_root, src_root):
        verilator_root = utils.get_verilator_root()
        if verilator_root is None:
            verilator_root = utils.get_verilator_root()
        args = ['-c']
        args += ['-I'+src_root]
        args += ['-I'+os.path.join(src_root, core.sanitized_name, s) for s in self.include_dirs]
        args += ['-Iobj_dir']
        args += ['-I'+os.path.join(verilator_root,'include')]
        args += ['-I'+os.path.join(verilator_root,'include', 'vltstd')]
        for src_file in core.verilator.src_files:
            pr_info("Compiling " + src_file.name)
            l = utils.Launcher('g++', args + [os.path.join(src_root, core.sanitized_name, src_file.name)],
                         cwd=work_root,
                         stderr = open(os.path.join(work_root, 'g++.err.log'),'a'))
            if Config().verbose:
                pr_info("  C++ compilation working dir: " + work_root)
                pr_info("  C++ compilation command: g++ " + ' '.join(args) + ' ' + os.path.join(src_root, core.sanitized_name, src_file.name))
            l.run()

    def build_SysC(self, core, work_root, src_root):
        verilator_root = utils.get_verilator_root()
        args = ['-I.']
        args += ['-MMD']
        args += ['-I'+src_root]
        args += ['-I'+s for s in self.include_dirs]
        args += ['-Iobj_dir']
        args += ['-I'+os.path.join(verilator_root,'include')]
        args += ['-I'+os.path.join(verilator_root,'include', 'vltstd')]
        args += ['-DVL_PRINTF=printf']
        args += ['-DVM_TRACE=1']
        args += ['-DVM_COVERAGE=0']
        if os.getenv('SYSTEMC_INCLUDE'):
            args += ['-I'+os.getenv('SYSTEMC_INCLUDE')]
        if os.getenv('SYSTEMC'):
            args += ['-I'+os.path.join(os.getenv('SYSTEMC'),'include')]
        args += ['-Wno-deprecated']
        if os.getenv('SYSTEMC_CXX_FLAGS'):
             args += [os.getenv('SYSTEMC_CXX_FLAGS')]
        args += ['-c']
        args += ['-g']

        for src_file in core.verilator.src_files:
            pr_info("Compiling " + src_file.name)
            l = utils.Launcher('g++', args + [os.path.join(src_root, core.sanitized_name, src_file.name)],
                         cwd=work_root,
                         stderr = open(os.path.join(work_root, 'g++.err.log'),'a'))
            if Config().verbose:
                pr_info("  SystemC compilation working dir: " + work_root)
                pr_info("  SystemC compilation command: g++ " + ' '.join(args) + ' ' + os.path.join(src_root, core.sanitized_name, src_file.name))
            l.run()

    def run(self, args):
        if not self.fusesoc_cli_parser:
            self.plusarg = []
        super(Verilator, self).run(args)
        self.env = os.environ.copy()
        self.env['CORE_ROOT'] = os.path.abspath(self.system.core_root)
        self.env['SIM_ROOT'] = os.path.abspath(self.work_root)
        if self.fusesoc_cli_parser:
            _args = []
            for key, value in self.plusarg.items():
                _args += ['+{}={}'.format(key, value)]
        else:
            _args = args
        pr_info("Running simulation")
        utils.Launcher('./V' + self.top_module,
                       _args,
                       cwd=os.path.join(self.work_root, 'obj_dir'),
                       env = self.env).run()
<|endoftext|>


<commit_before>
def readFVCOM(file, varList=None, clipDims=False, noisy=False, globalAtts=False):
    """
    Read in the FVCOM results file and spit out numpy arrays for each of the
    variables specified in the varList list.

    Optionally specify a dict with keys whose names match the dimension names
    in the NetCDF file and whose values are strings specifying alternative
    ranges or lists of indices. For example, to extract the first hundred time
    steps, supply clipDims as:

        clipDims = {'time':'0:100'}

    To extract the first, 400th and 10,000th values of any array with nodes:

        clipDims = {'node':['0, 3999, 9999']}

    Any dimension not given in clipDims will be extracted in full.

    Specify globalAtts=True to extract global attributes.

    Parameters
    ----------

    file : str, list
        If a string, the full path to an FVCOM NetCDF output file. If a list,
        a series of files to be loaded. Data will be concatenated into a single
        dict.
    varList : list, optional
        List of variable names to be extracted. If omitted, all variables are
        returned.
    clipDims : dict, optional
        Dict whose keys are dimensions and whose values are a string of either
        a range (e.g. {'time':'0:100'}) or a list of individual indices (e.g.
        {'time':'[0, 1, 80, 100]'}). Slicing is supported (::5 for every fifth
        value) but it is not possible to extract data from the end of the array
        with a negative index (e.g. 0:-4).
    noisy : bool
        Set to True to enable verbose output.
    globalAtts : bool, optional
        Set to True to enable output of the global attributes (defaults to
        False).

    Returns
    -------

    FVCOM : dict
        Dict of data extracted from the NetCDF file. Keys are those given in
        varList and the data are stored as ndarrays.
    attributes : dict, optional
        If globAtt=True, returns the global attributes as a dict for each
        variable in varList. The key 'dims' contains the array dimensions (each
        variable contains the names of its dimensions) as well as the shape of
        the dimensions defined in the NetCDF file.

    """

    try:
        from netCDF4 import Dataset, MFDataset
    except ImportError:
        raise ImportError('Failed to load the NetCDF4 library')

    import sys

    # If we have a list, assume it's lots of files and load them all.
    if isinstance(file, list):
        rootgrp = MFDataset(file, 'r')
    else:
        rootgrp = Dataset(file, 'r')

    # Create a dict of the dimension names and their current sizes
    dims = {}
    for key, var in rootgrp.dimensions.iteritems():
        # Make the dimensions ranges so we can use them to extract all the
        # values.
        dims[key] = '0:' + str(len(var))

    # Compare the dimensions in the NetCDF file with those provided. If we've
    # been given a dict of dimensions which differs from those in the NetCDF
    # file, then use those.
    if clipDims is not False:
        commonKeys = set(dims).intersection(clipDims.keys())
        for k in commonKeys:
            dims[k] = clipDims[k]

    if noisy:
        print "File format: {}".format(rootgrp.file_format)

    if varList is None:
        varList = rootgrp.variables.iterkeys()

    FVCOM = {}

    # Add the dimensions to the output dict. Hope we won't have a variable
    # named 'global'...
    attributes = {}
    attributes['dims'] = dims

    for key, var in rootgrp.variables.iteritems():
        if noisy:
            print 'Found ' + key,
            sys.stdout.flush()

        if key in varList:
            vDims = rootgrp.variables[key].dimensions

            toExtract = [dims[d] for d in vDims]

            # I know, I know, eval() is evil.
            getData = 'rootgrp.variables[\'{}\']{}'.format(key,str(toExtract).replace('\'', ''))
            FVCOM[key] = eval(getData)

            # Add the units and dimensions for this variable
            attributes[key] = {}
            try:
                attributes[key]['units'] = rootgrp.variables[key].units
            except:
                pass

            try:
                attributes[key]['dims'] = rootgrp.variables[key].dimensions
            except:
                pass

            if noisy:
                if len(str(toExtract)) < 60:
                    print '(extracted {})'.format(str(toExtract).replace('\'', ''))
                else:
                    print '(extracted given indices)'

        elif noisy:
                print

    # Close the open file.
    rootgrp.close()

    if globalAtts:
        return FVCOM, attributes
    else:
        return FVCOM

def elems2nodes(elems, tri, nvert, noisy=False):
    """
    Calculate a nodal value based on the average value for the elements
    of which it a part. This necessarily involves an average, so the
    conversion from nodes2elems and elems2nodes is not necessarily
    reversible.

    Parameters
    ----------

    elems : ndarray
        Array of unstructured grid element values to move to the element
        nodes.
    tri : ndarray
        Array of shape (nelem, 3) comprising the list of connectivity
        for each element.
    nvert : int
        Number of nodes (vertices) in the unstructured grid.

    Returns
    -------

    nodes : ndarray
        Array of values at the grid nodes.

    """

    try:
        import numpy as np
    except ImportError:
        raise ImportError('NumPy not found')

    count = np.zeros(nvert, dtype=int)

    # Deal with 1D and 2D element arrays separately
    if np.ndim(elems) == 1:
        nodes = np.zeros(nvert)
        for i, indices in enumerate(tri):
            n0, n1, n2 = indices
            nodes[n0] = nodes[n0] + elems[i]
            nodes[n1] = nodes[n1] + elems[i]
            nodes[n2] = nodes[n2] + elems[i]
            count[n0] = count[n0] + 1
            count[n1] = count[n1] + 1
            count[n2] = count[n2] + 1

    elif np.ndim(elems) == 2:
        nodes = np.zeros((np.shape(elems)[0], nvert))
        for i, indices in enumerate(tri):
            n0, n1, n2 = indices
            nodes[:, n0] = nodes[:, n0] + elems[:, i]
            nodes[:, n1] = nodes[:, n1] + elems[:, i]
            nodes[:, n2] = nodes[:, n2] + elems[:, i]
            count[n0] = count[n0] + 1
            count[n1] = count[n1] + 1
            count[n2] = count[n2] + 1
    else:
        raise 'Too many dimensions (maximum of two)'

    # Now calculate the average for each node based on the number of
    # elements of which it is a part.
    nodes = nodes / count

    return nodes

def nodes2elems(nodes, tri, noisy=False):
    """
    Calculate a element centre value based on the average value for the
    nodes from which it is formed. This necessarily involves an average,
    so the conversion from nodes2elems and elems2nodes is not
    necessarily reversible.

    Parameters
    ----------

    nodes : ndarray
        Array of unstructured grid node values to move to the element
        centres.
    tri : ndarray
        Array of shape (nelem, 3) comprising the list of connectivity
        for each element.

    Returns
    -------

    elems : ndarray
        Array of values at the grid nodes.

    """

    nvert = np.shape(tri)[0]

    if np.ndim(elems) == 1:
        elems = np.zeros(nvert)
        for i, indices in enumerate(tri):
            elems[i] = np.mean(nodes[indices])

    elif np.ndim(elems) == 2:
        elems = np.zeros((np.shape(nodes)[0], nvert))
        for i, indices in enumerate(tri):
            elems[:, i] = np.mean(nodes[:, indices])
    else:
        raise 'Too many dimensions (maximum of two)'

def getSurfaceElevation(Z, idx):
    """
    Extract the surface elevation from Z at index ind. If ind is multiple
    values, extract and return the surface elevations at all those locations.

    Z is usually extracted from the dict created when using readFVCOM() on a
    NetCDF file.

    Parameters
    ----------

    Z : ndarray
        Unstructured array of surface elevations with time.
    idx : list
        List of indices from which to extract time series of surface
        elevations.

    Returns
    -------

    surfaceElevation : ndarray
        Time series of surface elevations at the indices supplied in
        idx.

    """
    try:
        import numpy as np
    except ImportError:
        raise ImportError('NumPy not found')

    nt, nx = np.shape(Z)

    surfaceElevation = np.empty([nt,np.shape(idx)[0]])
    for cnt, i in enumerate(idx):
        if not np.isnan(i):
            surfaceElevation[:,cnt] = Z[:,i]

    return surfaceElevation
<commit_msg>Try aggregating files along a time dimension if searching for an unlimited dimension has failed. This is useful for POLCOMS output files which do not specify an unlimited time dimension.<commit_after>
def readFVCOM(file, varList=None, clipDims=False, noisy=False, globalAtts=False):
    """
    Read in the FVCOM results file and spit out numpy arrays for each of the
    variables specified in the varList list.

    Optionally specify a dict with keys whose names match the dimension names
    in the NetCDF file and whose values are strings specifying alternative
    ranges or lists of indices. For example, to extract the first hundred time
    steps, supply clipDims as:

        clipDims = {'time':'0:100'}

    To extract the first, 400th and 10,000th values of any array with nodes:

        clipDims = {'node':['0, 3999, 9999']}

    Any dimension not given in clipDims will be extracted in full.

    Specify globalAtts=True to extract global attributes.

    Parameters
    ----------

    file : str, list
        If a string, the full path to an FVCOM NetCDF output file. If a list,
        a series of files to be loaded. Data will be concatenated into a single
        dict.
    varList : list, optional
        List of variable names to be extracted. If omitted, all variables are
        returned.
    clipDims : dict, optional
        Dict whose keys are dimensions and whose values are a string of either
        a range (e.g. {'time':'0:100'}) or a list of individual indices (e.g.
        {'time':'[0, 1, 80, 100]'}). Slicing is supported (::5 for every fifth
        value) but it is not possible to extract data from the end of the array
        with a negative index (e.g. 0:-4).
    noisy : bool
        Set to True to enable verbose output.
    globalAtts : bool, optional
        Set to True to enable output of the global attributes (defaults to
        False).

    Returns
    -------

    FVCOM : dict
        Dict of data extracted from the NetCDF file. Keys are those given in
        varList and the data are stored as ndarrays.
    attributes : dict, optional
        If globAtt=True, returns the global attributes as a dict for each
        variable in varList. The key 'dims' contains the array dimensions (each
        variable contains the names of its dimensions) as well as the shape of
        the dimensions defined in the NetCDF file.

    """

    try:
        from netCDF4 import Dataset, MFDataset
    except ImportError:
        raise ImportError('Failed to load the NetCDF4 library')

    import sys

    # If we have a list, assume it's lots of files and load them all.
    if isinstance(file, list):
        try:
            rootgrp = MFDataset(file, 'r')
        except:
            # Try aggregating along a 'time' dimension (for POLCOMS, for example)
            rootgrp = MFDataset(file, 'r', aggdim='time')

    else:
        rootgrp = Dataset(file, 'r')

    # Create a dict of the dimension names and their current sizes
    dims = {}
    for key, var in rootgrp.dimensions.iteritems():
        # Make the dimensions ranges so we can use them to extract all the
        # values.
        dims[key] = '0:' + str(len(var))

    # Compare the dimensions in the NetCDF file with those provided. If we've
    # been given a dict of dimensions which differs from those in the NetCDF
    # file, then use those.
    if clipDims is not False:
        commonKeys = set(dims).intersection(clipDims.keys())
        for k in commonKeys:
            dims[k] = clipDims[k]

    if noisy:
        print "File format: {}".format(rootgrp.file_format)

    if varList is None:
        varList = rootgrp.variables.iterkeys()

    FVCOM = {}

    # Add the dimensions to the output dict. Hope we won't have a variable
    # named 'global'...
    attributes = {}
    attributes['dims'] = dims

    for key, var in rootgrp.variables.iteritems():
        if noisy:
            print 'Found ' + key,
            sys.stdout.flush()

        if key in varList:
            vDims = rootgrp.variables[key].dimensions

            toExtract = [dims[d] for d in vDims]

            # I know, I know, eval() is evil.
            getData = 'rootgrp.variables[\'{}\']{}'.format(key,str(toExtract).replace('\'', ''))
            FVCOM[key] = eval(getData)

            # Add the units and dimensions for this variable
            attributes[key] = {}
            try:
                attributes[key]['units'] = rootgrp.variables[key].units
            except:
                pass

            try:
                attributes[key]['dims'] = rootgrp.variables[key].dimensions
            except:
                pass

            if noisy:
                if len(str(toExtract)) < 60:
                    print '(extracted {})'.format(str(toExtract).replace('\'', ''))
                else:
                    print '(extracted given indices)'

        elif noisy:
                print

    # Close the open file.
    rootgrp.close()

    if globalAtts:
        return FVCOM, attributes
    else:
        return FVCOM

def elems2nodes(elems, tri, nvert, noisy=False):
    """
    Calculate a nodal value based on the average value for the elements
    of which it a part. This necessarily involves an average, so the
    conversion from nodes2elems and elems2nodes is not necessarily
    reversible.

    Parameters
    ----------

    elems : ndarray
        Array of unstructured grid element values to move to the element
        nodes.
    tri : ndarray
        Array of shape (nelem, 3) comprising the list of connectivity
        for each element.
    nvert : int
        Number of nodes (vertices) in the unstructured grid.

    Returns
    -------

    nodes : ndarray
        Array of values at the grid nodes.

    """

    try:
        import numpy as np
    except ImportError:
        raise ImportError('NumPy not found')

    count = np.zeros(nvert, dtype=int)

    # Deal with 1D and 2D element arrays separately
    if np.ndim(elems) == 1:
        nodes = np.zeros(nvert)
        for i, indices in enumerate(tri):
            n0, n1, n2 = indices
            nodes[n0] = nodes[n0] + elems[i]
            nodes[n1] = nodes[n1] + elems[i]
            nodes[n2] = nodes[n2] + elems[i]
            count[n0] = count[n0] + 1
            count[n1] = count[n1] + 1
            count[n2] = count[n2] + 1

    elif np.ndim(elems) == 2:
        nodes = np.zeros((np.shape(elems)[0], nvert))
        for i, indices in enumerate(tri):
            n0, n1, n2 = indices
            nodes[:, n0] = nodes[:, n0] + elems[:, i]
            nodes[:, n1] = nodes[:, n1] + elems[:, i]
            nodes[:, n2] = nodes[:, n2] + elems[:, i]
            count[n0] = count[n0] + 1
            count[n1] = count[n1] + 1
            count[n2] = count[n2] + 1
    else:
        raise 'Too many dimensions (maximum of two)'

    # Now calculate the average for each node based on the number of
    # elements of which it is a part.
    nodes = nodes / count

    return nodes

def nodes2elems(nodes, tri, noisy=False):
    """
    Calculate a element centre value based on the average value for the
    nodes from which it is formed. This necessarily involves an average,
    so the conversion from nodes2elems and elems2nodes is not
    necessarily reversible.

    Parameters
    ----------

    nodes : ndarray
        Array of unstructured grid node values to move to the element
        centres.
    tri : ndarray
        Array of shape (nelem, 3) comprising the list of connectivity
        for each element.

    Returns
    -------

    elems : ndarray
        Array of values at the grid nodes.

    """

    nvert = np.shape(tri)[0]

    if np.ndim(elems) == 1:
        elems = np.zeros(nvert)
        for i, indices in enumerate(tri):
            elems[i] = np.mean(nodes[indices])

    elif np.ndim(elems) == 2:
        elems = np.zeros((np.shape(nodes)[0], nvert))
        for i, indices in enumerate(tri):
            elems[:, i] = np.mean(nodes[:, indices])
    else:
        raise 'Too many dimensions (maximum of two)'

def getSurfaceElevation(Z, idx):
    """
    Extract the surface elevation from Z at index ind. If ind is multiple
    values, extract and return the surface elevations at all those locations.

    Z is usually extracted from the dict created when using readFVCOM() on a
    NetCDF file.

    Parameters
    ----------

    Z : ndarray
        Unstructured array of surface elevations with time.
    idx : list
        List of indices from which to extract time series of surface
        elevations.

    Returns
    -------

    surfaceElevation : ndarray
        Time series of surface elevations at the indices supplied in
        idx.

    """
    try:
        import numpy as np
    except ImportError:
        raise ImportError('NumPy not found')

    nt, nx = np.shape(Z)

    surfaceElevation = np.empty([nt,np.shape(idx)[0]])
    for cnt, i in enumerate(idx):
        if not np.isnan(i):
            surfaceElevation[:,cnt] = Z[:,i]

    return surfaceElevation
<|endoftext|>


<commit_before>"""
Tools for handling items on the canvas.

Although Gaphas has quite a few useful tools, some tools need to be extended:
 - PlacementTool: should perform undo
 - HandleTool: should support adapter based connection protocol
 - TextEditTool: should support adapter based edit protocol
"""

import logging

from gaphas.aspect import Connector as ConnectorAspect
from gaphas.aspect import InMotion as InMotionAspect
from gaphas.aspect import ItemConnector
from gaphas.guide import GuidedItemInMotion
from gaphas.tool import ConnectHandleTool, HoverTool, ItemTool
from gaphas.tool import PlacementTool as _PlacementTool
from gaphas.tool import RubberbandTool, Tool, ToolChain
from gi.repository import Gdk

from gaphor.core import Transaction, transactional
from gaphor.diagram.connectors import Connector
from gaphor.diagram.event import DiagramItemPlaced
from gaphor.diagram.grouping import Group
from gaphor.diagram.inlineeditors import InlineEditor
from gaphor.diagram.presentation import Presentation

log = logging.getLogger(__name__)

@ConnectorAspect.register(Presentation)
class DiagramItemConnector(ItemConnector):
    """
    Handle Tool (acts on item handles) that uses the Connector protocol
    to connect items to one-another.

    It also adds handles to lines when a line is grabbed on the middle of
    a line segment (points are drawn by the LineSegmentPainter).
    """

    def allow(self, sink):
        adapter = Connector(sink.item, self.item)
        return adapter and adapter.allow(self.handle, sink.port)

    @transactional
    def connect(self, sink):
        """
        Create connection at handle level and at model level.
        """
        handle = self.handle
        item = self.item
        cinfo = item.canvas.get_connection(handle)

        try:
            callback = DisconnectHandle(self.item, self.handle)
            if cinfo and cinfo.connected is sink.item:
                # reconnect only constraint - leave model intact
                log.debug("performing reconnect constraint")
                constraint = sink.port.constraint(item.canvas, item, handle, sink.item)
                item.canvas.reconnect_item(
                    item, handle, sink.port, constraint=constraint
                )
            elif cinfo:
                # first disconnect but disable disconnection handle as
                # reconnection is going to happen
                adapter = Connector(sink.item, item)
                try:
                    connect = adapter.reconnect
                except AttributeError:
                    connect = adapter.connect
                else:
                    cinfo.callback.disable = True
                self.disconnect()

                # new connection
                self.connect_handle(sink, callback=callback)

                # adapter requires both ends to be connected.
                connect(handle, sink.port)
            else:
                # new connection
                adapter = Connector(sink.item, item)
                self.connect_handle(sink, callback=callback)
                adapter.connect(handle, sink.port)
        except Exception:
            log.error("Error during connect", exc_info=True)

    @transactional
    def disconnect(self):
        super().disconnect()

class DisconnectHandle:
    """
    Callback for items disconnection using the adapters.

    This is an object so disconnection data can be serialized/deserialized
    using pickle.

    :Variables:
     item
        Connecting item.
     handle
        Handle of connecting item.
     disable
        If set, then disconnection is disabled.
    """

    def __init__(self, item, handle):
        self.item = item
        self.handle = handle
        self.disable = False

    def __call__(self):
        handle = self.handle
        item = self.item
        canvas = self.item.canvas
        cinfo = canvas.get_connection(handle)

        if self.disable:
            log.debug(f"Not disconnecting {item}.{handle} (disabled)")
        else:
            log.debug(f"Disconnecting {item}.{handle}")
            if cinfo:
                adapter = Connector(cinfo.connected, item)
                adapter.disconnect(handle)

class TextEditTool(Tool):
    """
    Text edit tool. Allows for elements that can adapt to the
    IEditable interface to be edited.
    """

    def on_key_press(self, event):
        view = self.view
        item = view.hovered_item
        if (
            item
            and event.type == Gdk.EventType.KEY_PRESS
            and event.key.keyval == Gdk.KEY_F2
        ):
            return InlineEditor(item, view)
        return False

    def on_double_click(self, event):
        view = self.view
        item = view.hovered_item
        if item:
            x, y = view.get_matrix_v2i(item).transform_point(event.x, event.y)
            return InlineEditor(item, view, (x, y))
        return False

class PlacementTool(_PlacementTool):
    """
    PlacementTool is used to place items on the canvas.
    """

    def __init__(self, view, item_factory, event_manager, handle_index=-1):
        """
        item_factory is a callable. It is used to create a CanvasItem
        that is displayed on the diagram.
        """
        _PlacementTool.__init__(
            self,
            view,
            factory=item_factory,
            handle_tool=ConnectHandleTool(),
            handle_index=handle_index,
        )
        self.event_manager = event_manager
        self._parent = None

    @classmethod
    def new_item_factory(_class, item_class, subject_class=None, config_func=None):
        """
        ``config_func`` may be a function accepting the newly created item.
        """

        def item_factory(diagram, parent=None):
            if subject_class:
                element_factory = diagram.model
                subject = element_factory.create(subject_class)
            else:
                subject = None

            item = diagram.create(item_class, subject=subject)
            if config_func:
                config_func(item)

            adapter = Group(parent, item)
            if parent and adapter.can_contain():
                canvas = diagram.canvas
                canvas.reparent(item, parent=parent)
                m = canvas.get_matrix_c2i(parent)
                item.matrix *= m
                adapter.group()

            return item

        item_factory.item_class = item_class  # type: ignore[attr-defined] # noqa: F821
        return item_factory

    @transactional
    def create_item(self, pos):
        """
        Create an item directly.
        """
        return self._create_item(pos)

    def on_button_press(self, event):
        view = self.view
        view.unselect_all()
        if super().on_button_press(event):
            try:
                opposite = self.new_item.opposite(
                    self.new_item.handles()[self._handle_index]
                )
            except (KeyError, AttributeError):
                pass
            else:
                # Connect opposite handle first, using the HandleTool's
                # mechanisms

                # First make sure all matrices are updated:
                view.canvas.update_matrices([self.new_item])
                view.update_matrix(self.new_item)

                vpos = event.x, event.y

                item = self.handle_tool.glue(self.new_item, opposite, vpos)
                if item:
                    self.handle_tool.connect(self.new_item, opposite, vpos)
            return True
        return False

    def on_button_release(self, event):
        self.event_manager.handle(DiagramItemPlaced(self.new_item))
        return super().on_button_release(event)

    def on_motion_notify(self, event):
        """
        Change parent item to dropzone state if it can accept diagram item
        object to be created.
        """
        if self.grabbed_handle:
            return self.handle_tool.on_motion_notify(event)

        view = self.view

        try:
            parent = view.get_item_at_point((event.x, event.y))
        except KeyError:
            parent = None

        if parent:
            # create dummy adapter
            adapter = Group(parent, self._factory.item_class())
            if adapter and adapter.can_contain():
                view.dropzone_item = parent
                self._parent = parent
            else:
                view.dropzone_item = None
                self._parent = None
            parent.request_update(matrix=False)
        else:
            if view.dropzone_item:
                view.dropzone_item.request_update(matrix=False)
            view.dropzone_item = None

    def _create_item(self, pos):
        """
        Create diagram item and place it within parent's boundaries.
        """
        parent = self._parent
        view = self.view
        diagram = view.canvas.diagram
        try:
            item = super()._create_item(pos, diagram=diagram, parent=parent)
        finally:
            self._parent = None
            view.dropzone_item = None
        return item

@InMotionAspect.register(Presentation)
class DropZoneInMotion(GuidedItemInMotion):
    def move(self, pos):
        """
        Move the item. x and y are in view coordinates.
        """
        super().move(pos)
        item = self.item
        view = self.view
        x, y = pos

        current_parent = view.canvas.get_parent(item)
        over_item = view.get_item_at_point((x, y), selected=False)

        if not over_item:
            view.dropzone_item = None
            return

        if current_parent and not over_item:
            # are we going to remove from parent?
            group = Group(current_parent, item)
            if group:
                view.dropzone_item = current_parent
                current_parent.request_update(matrix=False)

        if over_item:
            # are we going to add to parent?
            group = Group(over_item, item)
            if group and group.can_contain():
                view.dropzone_item = over_item
                over_item.request_update(matrix=False)

    def stop_move(self):
        """
        Motion stops: drop!
        """
        super().stop_move()
        item = self.item
        view = self.view
        canvas = view.canvas
        old_parent = view.canvas.get_parent(item)
        new_parent = view.dropzone_item
        try:

            if new_parent is old_parent:
                if old_parent is not None:
                    old_parent.request_update(matrix=False)
                return

            if old_parent:
                canvas.reparent(item, None)
                m = canvas.get_matrix_i2c(old_parent)
                item.matrix *= m

                adapter = Group(old_parent, item)
                if adapter:
                    adapter.ungroup()

                old_parent.request_update()

            if new_parent:
                canvas.reparent(item, new_parent)
                m = canvas.get_matrix_c2i(new_parent)
                item.matrix *= m

                adapter = Group(new_parent, item)
                if adapter and adapter.can_contain():
                    adapter.group()

                new_parent.request_update()
        finally:
            view.dropzone_item = None

class TransactionalToolChain(ToolChain):
    """
    In addition to a normal toolchain, this chain begins an undo-transaction
    at button-press and commits the transaction at button-release.
    """

    def __init__(self, event_manager, view=None):
        super().__init__(view)
        self.event_manager = event_manager
        self._tx = None

    def handle(self, event):
        # For double click: button_press, double_click, button_release
        if self.EVENT_HANDLERS.get(event.type) in ("on_button_press",):
            assert not self._tx
            self._tx = Transaction(self.event_manager)

        try:
            return super().handle(event)
        finally:
            if self._tx and self.EVENT_HANDLERS.get(event.type) in (
                "on_button_release",
                "on_double_click",
                "on_triple_click",
            ):
                self._tx.commit()
                self._tx = None

def DefaultTool(event_manager):
    """
    The default tool chain build from HoverTool, ItemTool and HandleTool.
    """
    chain = TransactionalToolChain(event_manager)
    chain.append(HoverTool())
    chain.append(ConnectHandleTool())
    chain.append(ItemTool())
    chain.append(TextEditTool())
    chain.append(RubberbandTool())
    return chain
<commit_msg>Fix grouping on placement<commit_after>"""
Tools for handling items on the canvas.

Although Gaphas has quite a few useful tools, some tools need to be extended:
 - PlacementTool: should perform undo
 - HandleTool: should support adapter based connection protocol
 - TextEditTool: should support adapter based edit protocol
"""

import logging

from gaphas.aspect import Connector as ConnectorAspect
from gaphas.aspect import InMotion as InMotionAspect
from gaphas.aspect import ItemConnector
from gaphas.guide import GuidedItemInMotion
from gaphas.tool import ConnectHandleTool, HoverTool, ItemTool
from gaphas.tool import PlacementTool as _PlacementTool
from gaphas.tool import RubberbandTool, Tool, ToolChain
from gi.repository import Gdk

from gaphor.core import Transaction, transactional
from gaphor.diagram.connectors import Connector
from gaphor.diagram.event import DiagramItemPlaced
from gaphor.diagram.grouping import Group
from gaphor.diagram.inlineeditors import InlineEditor
from gaphor.diagram.presentation import Presentation

log = logging.getLogger(__name__)

@ConnectorAspect.register(Presentation)
class DiagramItemConnector(ItemConnector):
    """
    Handle Tool (acts on item handles) that uses the Connector protocol
    to connect items to one-another.

    It also adds handles to lines when a line is grabbed on the middle of
    a line segment (points are drawn by the LineSegmentPainter).
    """

    def allow(self, sink):
        adapter = Connector(sink.item, self.item)
        return adapter and adapter.allow(self.handle, sink.port)

    @transactional
    def connect(self, sink):
        """
        Create connection at handle level and at model level.
        """
        handle = self.handle
        item = self.item
        cinfo = item.canvas.get_connection(handle)

        try:
            callback = DisconnectHandle(self.item, self.handle)
            if cinfo and cinfo.connected is sink.item:
                # reconnect only constraint - leave model intact
                log.debug("performing reconnect constraint")
                constraint = sink.port.constraint(item.canvas, item, handle, sink.item)
                item.canvas.reconnect_item(
                    item, handle, sink.port, constraint=constraint
                )
            elif cinfo:
                # first disconnect but disable disconnection handle as
                # reconnection is going to happen
                adapter = Connector(sink.item, item)
                try:
                    connect = adapter.reconnect
                except AttributeError:
                    connect = adapter.connect
                else:
                    cinfo.callback.disable = True
                self.disconnect()

                # new connection
                self.connect_handle(sink, callback=callback)

                # adapter requires both ends to be connected.
                connect(handle, sink.port)
            else:
                # new connection
                adapter = Connector(sink.item, item)
                self.connect_handle(sink, callback=callback)
                adapter.connect(handle, sink.port)
        except Exception:
            log.error("Error during connect", exc_info=True)

    @transactional
    def disconnect(self):
        super().disconnect()

class DisconnectHandle:
    """
    Callback for items disconnection using the adapters.

    This is an object so disconnection data can be serialized/deserialized
    using pickle.

    :Variables:
     item
        Connecting item.
     handle
        Handle of connecting item.
     disable
        If set, then disconnection is disabled.
    """

    def __init__(self, item, handle):
        self.item = item
        self.handle = handle
        self.disable = False

    def __call__(self):
        handle = self.handle
        item = self.item
        canvas = self.item.canvas
        cinfo = canvas.get_connection(handle)

        if self.disable:
            log.debug(f"Not disconnecting {item}.{handle} (disabled)")
        else:
            log.debug(f"Disconnecting {item}.{handle}")
            if cinfo:
                adapter = Connector(cinfo.connected, item)
                adapter.disconnect(handle)

class TextEditTool(Tool):
    """
    Text edit tool. Allows for elements that can adapt to the
    IEditable interface to be edited.
    """

    def on_key_press(self, event):
        view = self.view
        item = view.hovered_item
        if (
            item
            and event.type == Gdk.EventType.KEY_PRESS
            and event.key.keyval == Gdk.KEY_F2
        ):
            return InlineEditor(item, view)
        return False

    def on_double_click(self, event):
        view = self.view
        item = view.hovered_item
        if item:
            x, y = view.get_matrix_v2i(item).transform_point(event.x, event.y)
            return InlineEditor(item, view, (x, y))
        return False

class PlacementTool(_PlacementTool):
    """
    PlacementTool is used to place items on the canvas.
    """

    def __init__(self, view, item_factory, event_manager, handle_index=-1):
        """
        item_factory is a callable. It is used to create a CanvasItem
        that is displayed on the diagram.
        """
        _PlacementTool.__init__(
            self,
            view,
            factory=item_factory,
            handle_tool=ConnectHandleTool(),
            handle_index=handle_index,
        )
        self.event_manager = event_manager
        self._parent = None

    @classmethod
    def new_item_factory(_class, item_class, subject_class=None, config_func=None):
        """
        ``config_func`` may be a function accepting the newly created item.
        """

        def item_factory(diagram, parent=None):
            if subject_class:
                element_factory = diagram.model
                subject = element_factory.create(subject_class)
            else:
                subject = None

            item = diagram.create(item_class, subject=subject)
            if config_func:
                config_func(item)

            adapter = Group(parent, item)
            if parent and adapter.can_contain():
                canvas = diagram.canvas
                canvas.reparent(item, parent=parent)
                adapter.group()

            return item

        item_factory.item_class = item_class  # type: ignore[attr-defined] # noqa: F821
        return item_factory

    @transactional
    def create_item(self, pos):
        """
        Create an item directly.
        """
        return self._create_item(pos)

    def on_button_press(self, event):
        view = self.view
        view.unselect_all()
        if super().on_button_press(event):
            try:
                opposite = self.new_item.opposite(
                    self.new_item.handles()[self._handle_index]
                )
            except (KeyError, AttributeError):
                pass
            else:
                # Connect opposite handle first, using the HandleTool's
                # mechanisms

                # First make sure all matrices are updated:
                view.canvas.update_matrices([self.new_item])
                view.update_matrix(self.new_item)

                vpos = event.x, event.y

                item = self.handle_tool.glue(self.new_item, opposite, vpos)
                if item:
                    self.handle_tool.connect(self.new_item, opposite, vpos)
            return True
        return False

    def on_button_release(self, event):
        self.event_manager.handle(DiagramItemPlaced(self.new_item))
        return super().on_button_release(event)

    def on_motion_notify(self, event):
        """
        Change parent item to dropzone state if it can accept diagram item
        object to be created.
        """
        if self.grabbed_handle:
            return self.handle_tool.on_motion_notify(event)

        view = self.view

        try:
            parent = view.get_item_at_point((event.x, event.y))
        except KeyError:
            parent = None

        if parent:
            # create dummy adapter
            adapter = Group(parent, self._factory.item_class())
            if adapter and adapter.can_contain():
                view.dropzone_item = parent
                self._parent = parent
            else:
                view.dropzone_item = None
                self._parent = None
            parent.request_update(matrix=False)
        else:
            if view.dropzone_item:
                view.dropzone_item.request_update(matrix=False)
            view.dropzone_item = None

    def _create_item(self, pos):
        """
        Create diagram item and place it within parent's boundaries.
        """
        parent = self._parent
        view = self.view
        diagram = view.canvas.diagram
        try:
            item = super()._create_item(pos, diagram=diagram, parent=parent)
        finally:
            self._parent = None
            view.dropzone_item = None
        return item

@InMotionAspect.register(Presentation)
class DropZoneInMotion(GuidedItemInMotion):
    def move(self, pos):
        """
        Move the item. x and y are in view coordinates.
        """
        super().move(pos)
        item = self.item
        view = self.view
        x, y = pos

        current_parent = view.canvas.get_parent(item)
        over_item = view.get_item_at_point((x, y), selected=False)

        if not over_item:
            view.dropzone_item = None
            return

        if current_parent and not over_item:
            # are we going to remove from parent?
            group = Group(current_parent, item)
            if group:
                view.dropzone_item = current_parent
                current_parent.request_update(matrix=False)

        if over_item:
            # are we going to add to parent?
            group = Group(over_item, item)
            if group and group.can_contain():
                view.dropzone_item = over_item
                over_item.request_update(matrix=False)

    def stop_move(self):
        """
        Motion stops: drop!
        """
        super().stop_move()
        item = self.item
        view = self.view
        canvas = view.canvas
        old_parent = view.canvas.get_parent(item)
        new_parent = view.dropzone_item
        try:

            if new_parent is old_parent:
                if old_parent is not None:
                    old_parent.request_update(matrix=False)
                return

            if old_parent:
                canvas.reparent(item, None)

                adapter = Group(old_parent, item)
                if adapter:
                    adapter.ungroup()

                old_parent.request_update()

            if new_parent:
                canvas.reparent(item, new_parent)

                adapter = Group(new_parent, item)
                if adapter and adapter.can_contain():
                    adapter.group()

                new_parent.request_update()
        finally:
            view.dropzone_item = None

class TransactionalToolChain(ToolChain):
    """
    In addition to a normal toolchain, this chain begins an undo-transaction
    at button-press and commits the transaction at button-release.
    """

    def __init__(self, event_manager, view=None):
        super().__init__(view)
        self.event_manager = event_manager
        self._tx = None

    def handle(self, event):
        # For double click: button_press, double_click, button_release
        if self.EVENT_HANDLERS.get(event.type) in ("on_button_press",):
            assert not self._tx
            self._tx = Transaction(self.event_manager)

        try:
            return super().handle(event)
        finally:
            if self._tx and self.EVENT_HANDLERS.get(event.type) in (
                "on_button_release",
                "on_double_click",
                "on_triple_click",
            ):
                self._tx.commit()
                self._tx = None

def DefaultTool(event_manager):
    """
    The default tool chain build from HoverTool, ItemTool and HandleTool.
    """
    chain = TransactionalToolChain(event_manager)
    chain.append(HoverTool())
    chain.append(ConnectHandleTool())
    chain.append(ItemTool())
    chain.append(TextEditTool())
    chain.append(RubberbandTool())
    return chain
<|endoftext|>


<commit_before>"""
Utilities for UConnRCMPy
"""
# System imports
from datetime import datetime

# Third Party imports
import win32clipboard

class ParsedFilename(object):
    """
    Parse a filename for experimental conditions.
    """

    def __init__(self, file_path):
        """
        Given a filename of the form
        `[NR_]XX_in_YY_mm_ZZZK-AAAAt-BBBx-DD-Mon-YY-Time.txt,`
        parse the filename for the experimental conditions.

        * `[NR_] = Optional non-reactive indicator
        * `XX` = inches of spacers
        * `YY` = millimeters of shims
        * `ZZZ` = Initial temperature in Kelvin
        * `AAAA` = Initial pressure in Torr
        * `BBB` = Multiplication factor set on the charge amplifier
        * `DD-Mon-YY-Time` = Day, Month, year, and time of experiment
        """
        fname = file_path.name.lstrip('NR_')
        """Processed filename."""
        fname = fname.rstrip('.txt')
        self.name_split = fname.split('_')
        """Split list of filename parts."""
        self.spacers = int(self.name_split[0])/10
        """Inches of spacers."""
        self.shims = int(self.name_split[2])
        """Millimeters of shims."""
        self.name_split_end = self.name_split[4].split('-', maxsplit=3)
        """Split list of filename parts."""
        self.Tin = int(self.name_split_end[0][:-1])
        """Initial temperature in Kelvin."""
        self.pin = int(self.name_split_end[1][:-1])
        """Initial pressure in torr."""
        self.factor = int(self.name_split_end[2][:-1])
        """Multiplication factor from charge amplifier."""
        self.data_date = datetime.strptime(self.name_split_end[3],
                                           '%d-%b-%y-%H%M')
        """Datetime object storing the processed date."""
        self.time_of_day = self.data_date.strftime('%H%M')
        """Time of day of the experiment."""
        self.date = self.data_date.strftime('%d-%b-%H%M')
        """Date of the experiment."""

def copy(text):
    """Copy the input `text` to the Windows clipboard."""
    win32clipboard.OpenClipboard()
    win32clipboard.EmptyClipboard()
    win32clipboard.SetClipboardText(text)
    win32clipboard.CloseClipboard()
<commit_msg>Add parse_file_name function<commit_after>"""
Utilities for UConnRCMPy
"""
# System imports
from datetime import datetime

# Third Party imports
import win32clipboard

def parse_file_name(file_path):
    """Parse the file name of an experimental trace.

    Parameters
    ----------
    file_path : Path
        The Path object that contains the path of the current experimental data file.
        The filename associated with the path should be in the following format::

            [NR_]XX_in_YY_mm_ZZZK-AAAAt-BBBx-DD-Mon-YY-Time.txt

        where:

        - `[NR_]`: Optional non-reactive indicator
        - `XX`: inches of spacers
        - `YY`: millimeters of shims
        - `ZZZ`: Initial temperature in Kelvin
        - `AAAA`: Initial pressure in Torr
        - `BBB`: Multiplication factor set on the charge amplifier
        - `DD-Mon-YY-Time`: Day, Month, Year, and Time of experiment

    Returns
    -------
    dict
        Dictionary containing the parameters of the experiment with the
        following names::

        - spacers: Inches of spacers
        - shims: Millimeters of shims
        - Tin: Initial temperature in Kelvin
        - pin: Initial pressure in Torr
        - factor: Multiplication factor set on the charge amplifier
        - time_of_day: Time of day of the experiment
        - date: Date of the experiment

    """
    name_parts = {}
    fname = file_path.name.lstrip('NR_')
    fname = fname.rstrip('.txt')
    name_split = fname.split('_')
    name_split_end = name_split[4].split('-', maxsplit=3)
    data_date = datetime.strptime(name_split_end[3], '%d-%b-%y-%H%M')

    name_parts['spacers'] = int(name_split[0])/10
    name_parts['shims'] = int(name_split[2])
    name_parts['Tin'] = int(name_split_end[0][:-1])
    name_parts['pin'] = int(name_split_end[1][:-1])
    name_parts['factor'] = int(name_split_end[2][:-1])
    name_parts['time_of_day'] = data_date.strftime('%H%M')
    name_parts['date'] = data_date.strftime('%d-%b-%H%M')
    return name_parts

class ParsedFilename(object):
    """
    Parse a filename for experimental conditions.
    """

    def __init__(self, file_path):
        """
        Given a filename of the form
        `[NR_]XX_in_YY_mm_ZZZK-AAAAt-BBBx-DD-Mon-YY-Time.txt,`
        parse the filename for the experimental conditions.

        * `[NR_] = Optional non-reactive indicator
        * `XX` = inches of spacers
        * `YY` = millimeters of shims
        * `ZZZ` = Initial temperature in Kelvin
        * `AAAA` = Initial pressure in Torr
        * `BBB` = Multiplication factor set on the charge amplifier
        * `DD-Mon-YY-Time` = Day, Month, year, and time of experiment
        """
        fname = file_path.name.lstrip('NR_')
        """Processed filename."""
        fname = fname.rstrip('.txt')
        self.name_split = fname.split('_')
        """Split list of filename parts."""
        self.spacers = int(self.name_split[0])/10
        """Inches of spacers."""
        self.shims = int(self.name_split[2])
        """Millimeters of shims."""
        self.name_split_end = self.name_split[4].split('-', maxsplit=3)
        """Split list of filename parts."""
        self.Tin = int(self.name_split_end[0][:-1])
        """Initial temperature in Kelvin."""
        self.pin = int(self.name_split_end[1][:-1])
        """Initial pressure in torr."""
        self.factor = int(self.name_split_end[2][:-1])
        """Multiplication factor from charge amplifier."""
        self.data_date = datetime.strptime(self.name_split_end[3],
                                           '%d-%b-%y-%H%M')
        """Datetime object storing the processed date."""
        self.time_of_day = self.data_date.strftime('%H%M')
        """Time of day of the experiment."""
        self.date = self.data_date.strftime('%d-%b-%H%M')
        """Date of the experiment."""

def copy(text):
    """Copy the input `text` to the Windows clipboard."""
    win32clipboard.OpenClipboard()
    win32clipboard.EmptyClipboard()
    win32clipboard.SetClipboardText(text)
    win32clipboard.CloseClipboard()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import math

def kageInt(s):
    """Imitates Math.floor in ECMAScript (returns an int).

    KAGE Engine uses Math.floor to parse numbers in KAGE data.
    """
    if isinstance(s, str):
        s = s.strip()
        if s == "":
            return 0
    elif isinstance(s, float):
        return int(math.floor(s))
    try:
        # decimal integer literal (may have leading 0 digits)
        return int(s)
    except ValueError:
        try:
            # hexadecimal, octal, binary integer literal
            return int(s, 0)
        except ValueError:
            # contains decimal point and/or exponent part
            return int(math.floor(float(s)))

def kageIntSuppressError(s):
    """The same as kageInt except that it returns None when s is invalid"""
    try:
        return kageInt(s)
    except (ValueError, OverflowError):
        return None

class KageData(object):

    def __init__(self, data):
        self.lines = tuple([KageLine(i, l)
                            for i, l in enumerate(data.split("$"))])
        self.len = len(self.lines)

    def isAlias(self):
        return self.len == 1 and self.lines[0].strdata[:19] == "99:0:0:0:0:200:200:"

class KageLine(object):

    def __init__(self, line_number, data):
        self.line_number = line_number
        self.strdata = data
        sdata = data.split(":")
        if kageIntSuppressError(sdata[0]) != 99:
            self.data = tuple([kageIntSuppressError(x) for x in sdata])
        else:
            self.data = tuple(
                [kageIntSuppressError(x) if i != 7 else x for i, x in enumerate(sdata)])
<commit_msg>Cache the result of isAlias<commit_after># -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import math

def kageInt(s):
    """Imitates Math.floor in ECMAScript (returns an int).

    KAGE Engine uses Math.floor to parse numbers in KAGE data.
    """
    if isinstance(s, str):
        s = s.strip()
        if s == "":
            return 0
    elif isinstance(s, float):
        return int(math.floor(s))
    try:
        # decimal integer literal (may have leading 0 digits)
        return int(s)
    except ValueError:
        try:
            # hexadecimal, octal, binary integer literal
            return int(s, 0)
        except ValueError:
            # contains decimal point and/or exponent part
            return int(math.floor(float(s)))

def kageIntSuppressError(s):
    """The same as kageInt except that it returns None when s is invalid"""
    try:
        return kageInt(s)
    except (ValueError, OverflowError):
        return None

class KageData(object):

    def __init__(self, data):
        self.lines = tuple([KageLine(i, l)
                            for i, l in enumerate(data.split("$"))])
        self.len = len(self.lines)

    def isAlias(self):
        res = self.len == 1 and self.lines[0].strdata[:19] == "99:0:0:0:0:200:200:"
        self.isAlias = lambda: res
        return res

class KageLine(object):

    def __init__(self, line_number, data):
        self.line_number = line_number
        self.strdata = data
        sdata = data.split(":")
        if kageIntSuppressError(sdata[0]) != 99:
            self.data = tuple([kageIntSuppressError(x) for x in sdata])
        else:
            self.data = tuple(
                [kageIntSuppressError(x) if i != 7 else x for i, x in enumerate(sdata)])
<|endoftext|>


<commit_before>import os
import gym
import ipdb
import math
import time
import random
import pickle
import argparse
import collections
import numpy as np
import tensorflow as tf

from cartpole_runner import *
from action_samplers import *
from models import N_DIM_ACTIONS, N_DIM_STATE, ddqn_mlp, ladder_mlp
from file_utils import *

config = tf.ConfigProto()
config.gpu_options.allow_growth=True
# config = tf.ConfigProto(gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.4))

print('#####################################################################################')
print('# Training')
print('#####################################################################################')
parser = argparse.ArgumentParser(description='Train an RL agent to learn Cart Pole')
parser.add_argument('--lr', help='learning rates to try', metavar='N', type=float, nargs='+')
parser.add_argument('--n_hidden', help='n hidden nodes', type=int)
parser.add_argument('--dir', help='output directory')
parser.add_argument('--use_ladder', dest='use_ladder', action='store_true')
parser.add_argument('--watch', dest='watch', action='store_true')
parser.add_argument('--experiments', help='To help batching experiments, comma separated list ' +\
        'of experiment indexes', type=str)
args = parser.parse_args()

# Training - Setup
DISCOUNT = 0.99
N_EPISODES = 1000
BATCH_SIZE = 16
N_STEPS_BETWEEN_TARGET_UPDATES = 1
N_EPISODES_PER_LOG = 1
EXPERIENCE_REPLAY_BUFFER_SIZE = 5000
learning_rates = [1e-3] if args.lr is None else args.lr
n_hidden = 100 if args.n_hidden is None else args.n_hidden
progress_dir = 'scratch' if args.dir is None else args.dir
use_ladder = args.use_ladder
watch = args.watch
progress_dir = os.path.join('models', progress_dir)
# env = gym.make('CartPole-v000')
env = gym.make('Enduro-v0')
start_time = time.time()
encoded_actions = []
for action_index in range(N_DIM_ACTIONS):
    a = np.zeros((N_DIM_ACTIONS), dtype=np.float32)
    a[action_index] = 1.0
    encoded_actions.append(a)
encoded_actions = np.array(encoded_actions)
experiment_indexes = [x for x in range(0, 25)]
if args.experiments is not None:
    experiment_indexes = [int(x) for x in args.experiments.split(',')]

print('Using learning rates: %s' % str(learning_rates))
print('Running experiments %s' % str(experiment_indexes))
print('Using network with %d hidden layers' % n_hidden)
print('Using ladder network: %s' % use_ladder)
print('#####################################################################################')

def state_converter(state):
    return np.mean(state, axis=(2)).flatten()

# Repeat the experiment for different learning rates
for learning_rate in learning_rates:
    print('Learning rate %f' % learning_rate)
    # Run the experiment 100 times to get a good training graph
    for curr_experiment in experiment_indexes:
        print('Running experiment %d... (%f s)' % (curr_experiment + 1, time.time() - start_time))
        test_lengths = np.zeros([int(N_EPISODES/N_EPISODES_PER_LOG)])
        lengths = np.zeros([int(N_EPISODES)])
        losses = np.zeros([int(N_EPISODES)])
        rewards = np.zeros([int(N_EPISODES)])
        experiences = collections.deque(maxlen=EXPERIENCE_REPLAY_BUFFER_SIZE)

        with tf.Session(config=config) as sess:
            s1_tf = tf.placeholder("float", [None, N_DIM_STATE], name='s1_inputs')
            a1_tf = tf.placeholder("float", [None, N_DIM_ACTIONS], name='a1_inputs')
            r1_tf = tf.placeholder("float", [None, 1], name='r1_inputs')
            s2_tf = tf.placeholder("float", [None, N_DIM_STATE], name='s2_inputs')
            # layers = [N_DIM_STATE, n_hidden, n_hidden, N_DIM_ACTIONS]
            layers = [N_DIM_STATE, 512, 128, 32, N_DIM_ACTIONS]
            denoising_cost = [10.0, 2.5, 0.0, 0.0, 0.0]
            if use_ladder:
                loss, train_op, best_action_picker, updater, training, tf_debug = \
                    ladder_mlp(s1_tf, a1_tf, r1_tf, s2_tf, DISCOUNT, learning_rate, layers, denoising_cost)
            else:
                loss, train_op, best_action_picker, updater, training, tf_debug = \
                    ddqn_mlp(s1_tf, a1_tf, r1_tf, s2_tf, DISCOUNT, learning_rate, layers)
            sess.run(tf.global_variables_initializer())
            action_sampler = e_greedy_sampler(sess, best_action_picker, s1_tf, training)
            saver = tf.train.Saver()

            # Train online Q-learning algorithm over N_EPISODES_PER_LOG00 episodes
            best_length = -float('inf')
            n_updates_since_updating_target_net = 0
            for curr_episode in range(0, N_EPISODES):
                # Log current status
                if (curr_episode+1) % N_EPISODES_PER_LOG == 0:
                    print('Exp %d - ep. %d: test length (%f) train length (%f) reward (%f) loss (%f) %f s' %
                          (curr_experiment + 1,
                           curr_episode + 1,
                           test_lengths[int(curr_episode / N_EPISODES_PER_LOG)],
                           np.mean(lengths[curr_episode - N_EPISODES_PER_LOG:curr_episode]),
                           rewards[curr_episode - 1],
                           losses[curr_episode - 1],
                           time.time() - start_time))

                # Run one episode
                state = env.reset()
                done = False
                curr_episode_length = 0
                curr_episode_reward = 0
                while not done:
                    action = action_sampler(env.action_space, np.array(state_converter(state), ndmin=2))
                    next_state, reward, done, _ = env.step(action)
                    if watch:
                        env.render()
                    # reward = -1 if done else 0
                    curr_episode_reward += reward
                    experiences.append((state_converter(state), encoded_actions[action, :], reward, state_converter(next_state)))
                    curr_episode_length += 1

                    # Train network from experiences saved in the replay buffer
                    exps = [experiences[random.randint(0, len(experiences)-1)] for i in range(0, BATCH_SIZE)]
                    sess.run(train_op, feed_dict={
                        s1_tf: np.array([s1 for s1, a1, r1, s2 in exps]),
                        a1_tf: np.array([a1 for s1, a1, r1, s2 in exps]),
                        r1_tf: np.array([r1 for s1, a1, r1, s2 in exps], ndmin=2).T,
                        s2_tf: np.array([s2 for s1, a1, r1, s2 in exps]),
                        training: True
                    })
                    state = next_state

                    # DDQN - Update the target network every N updates
                    n_updates_since_updating_target_net += 1
                    if n_updates_since_updating_target_net >= N_STEPS_BETWEEN_TARGET_UPDATES:
                        n_updates_since_updating_target_net = 0
                        updater(sess)

                lengths[curr_episode] = curr_episode_length
                rewards[curr_episode] = curr_episode_reward
                losses[curr_episode] = sess.run(loss, feed_dict={
                    s1_tf: np.array([s1 for s1, a1, r1, s2 in experiences]),
                    a1_tf: np.array([a1 for s1, a1, r1, s2 in experiences]),
                    r1_tf: np.array([r1 for s1, a1, r1, s2 in experiences], ndmin=2).T,
                    s2_tf: np.array([s2 for s1, a1, r1, s2 in experiences]),
                    training: True
                })
                if curr_episode % N_EPISODES_PER_LOG == 0:
                    # if tf_debug is not None:
                    #     print(sess.run(tf_debug, feed_dict={
                    #         s1_tf: np.array([s1 for s1, a1, r1, s2 in experiences]),
                    #         a1_tf: np.array([a1 for s1, a1, r1, s2 in experiences]),
                    #         r1_tf: np.array([r1 for s1, a1, r1, s2 in experiences], ndmin=2).T,
                    #         s2_tf: np.array([s2 for s1, a1, r1, s2 in experiences]),
                    #         training: True
                    #     }))
                    # _, episode_lengths, episode_rewards = \
                    #     run_n_cartpole_simulations(
                    #         25,
                    #         greedy_sampler(sess, best_action_picker, s1_tf, training),
                    #         env)
                    # test_lengths[int(curr_episode / N_EPISODES_PER_LOG)] = np.mean(np.array(episode_lengths))
                    # saved_progress_path = lr_filename(progress_dir, 'test_lengths_%d' % curr_experiment, learning_rate, 'npy')
                    # if not os.path.exists(os.path.dirname(saved_progress_path)):
                    #     os.makedirs(os.path.dirname(saved_progress_path))
                    # np.save(saved_progress_path, test_lengths)
                    saved_model_path = os.path.join(progress_dir, 'best_model_%d' % np.mean(np.array(rewards[curr_episode])))
                    if not os.path.exists(os.path.dirname(saved_model_path)):
                        os.makedirs(os.path.dirname(saved_model_path))
                    saver.save(sess, saved_model_path)

        np.save(lr_filename(progress_dir, 'lengths_%d' % curr_experiment, learning_rate, 'npy'), lengths)
        np.save(lr_filename(progress_dir, 'losses_%d' % curr_experiment, learning_rate, 'npy'), losses)
        np.save(lr_filename(progress_dir, 'rewards_%d' % curr_experiment, learning_rate, 'npy'), rewards)
        print('Saved! (%f)' % (time.time() - start_time))

<commit_msg>Switch to breakout<commit_after>import os
import gym
import ipdb
import math
import time
import random
import pickle
import argparse
import collections
import numpy as np
import tensorflow as tf

from cartpole_runner import *
from action_samplers import *
from models import N_DIM_ACTIONS, N_DIM_STATE, ddqn_mlp, ladder_mlp
from file_utils import *

config = tf.ConfigProto()
config.gpu_options.allow_growth=True
# config = tf.ConfigProto(gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.4))

print('#####################################################################################')
print('# Training')
print('#####################################################################################')
parser = argparse.ArgumentParser(description='Train an RL agent to learn Cart Pole')
parser.add_argument('--lr', help='learning rates to try', metavar='N', type=float, nargs='+')
parser.add_argument('--n_hidden', help='n hidden nodes', type=int)
parser.add_argument('--dir', help='output directory')
parser.add_argument('--use_ladder', dest='use_ladder', action='store_true')
parser.add_argument('--watch', dest='watch', action='store_true')
parser.add_argument('--experiments', help='To help batching experiments, comma separated list ' +\
        'of experiment indexes', type=str)
args = parser.parse_args()

# Training - Setup
DISCOUNT = 0.99
N_EPISODES = 1000
BATCH_SIZE = 16
N_STEPS_BETWEEN_TARGET_UPDATES = 1
N_EPISODES_PER_LOG = 1
EXPERIENCE_REPLAY_BUFFER_SIZE = 5000
learning_rates = [1e-3] if args.lr is None else args.lr
n_hidden = 100 if args.n_hidden is None else args.n_hidden
progress_dir = 'scratch' if args.dir is None else args.dir
use_ladder = args.use_ladder
watch = args.watch
progress_dir = os.path.join('models', progress_dir)
# env = gym.make('CartPole-v000')
env = gym.make('Breakout-v0')
start_time = time.time()
encoded_actions = []
for action_index in range(N_DIM_ACTIONS):
    a = np.zeros((N_DIM_ACTIONS), dtype=np.float32)
    a[action_index] = 1.0
    encoded_actions.append(a)
encoded_actions = np.array(encoded_actions)
experiment_indexes = [x for x in range(0, 25)]
if args.experiments is not None:
    experiment_indexes = [int(x) for x in args.experiments.split(',')]

print('Using learning rates: %s' % str(learning_rates))
print('Running experiments %s' % str(experiment_indexes))
print('Using network with %d hidden layers' % n_hidden)
print('Using ladder network: %s' % use_ladder)
print('#####################################################################################')

def state_converter(state):
    return np.mean(state, axis=(2)).flatten()

# Repeat the experiment for different learning rates
for learning_rate in learning_rates:
    print('Learning rate %f' % learning_rate)
    # Run the experiment 100 times to get a good training graph
    for curr_experiment in experiment_indexes:
        print('Running experiment %d... (%f s)' % (curr_experiment + 1, time.time() - start_time))
        test_lengths = np.zeros([int(N_EPISODES/N_EPISODES_PER_LOG)])
        lengths = np.zeros([int(N_EPISODES)])
        losses = np.zeros([int(N_EPISODES)])
        rewards = np.zeros([int(N_EPISODES)])
        experiences = collections.deque(maxlen=EXPERIENCE_REPLAY_BUFFER_SIZE)

        with tf.Session(config=config) as sess:
            s1_tf = tf.placeholder("float", [None, N_DIM_STATE], name='s1_inputs')
            a1_tf = tf.placeholder("float", [None, N_DIM_ACTIONS], name='a1_inputs')
            r1_tf = tf.placeholder("float", [None, 1], name='r1_inputs')
            s2_tf = tf.placeholder("float", [None, N_DIM_STATE], name='s2_inputs')
            # layers = [N_DIM_STATE, n_hidden, n_hidden, N_DIM_ACTIONS]
            layers = [N_DIM_STATE, 512, 128, 32, N_DIM_ACTIONS]
            denoising_cost = [10.0, 2.5, 0.0, 0.0, 0.0]
            if use_ladder:
                loss, train_op, best_action_picker, updater, training, tf_debug = \
                    ladder_mlp(s1_tf, a1_tf, r1_tf, s2_tf, DISCOUNT, learning_rate, layers, denoising_cost)
            else:
                loss, train_op, best_action_picker, updater, training, tf_debug = \
                    ddqn_mlp(s1_tf, a1_tf, r1_tf, s2_tf, DISCOUNT, learning_rate, layers)
            sess.run(tf.global_variables_initializer())
            action_sampler = e_greedy_sampler(sess, best_action_picker, s1_tf, training)
            saver = tf.train.Saver()

            # Train online Q-learning algorithm over N_EPISODES_PER_LOG00 episodes
            best_length = -float('inf')
            n_updates_since_updating_target_net = 0
            for curr_episode in range(0, N_EPISODES):
                # Log current status
                if (curr_episode+1) % N_EPISODES_PER_LOG == 0:
                    print('Exp %d - ep. %d: test length (%f) train length (%f) reward (%f) loss (%f) %f s' %
                          (curr_experiment + 1,
                           curr_episode + 1,
                           test_lengths[int(curr_episode / N_EPISODES_PER_LOG)],
                           np.mean(lengths[curr_episode - N_EPISODES_PER_LOG:curr_episode]),
                           rewards[curr_episode - 1],
                           losses[curr_episode - 1],
                           time.time() - start_time))

                # Run one episode
                state = env.reset()
                done = False
                curr_episode_length = 0
                curr_episode_reward = 0
                lives = 0
                while not done:
                    if curr_episode_length % 10 == 0:
                        print(curr_episode_length)
                    action = action_sampler(env.action_space, np.array(state_converter(state), ndmin=2))
                    next_state, reward, done, info = env.step(action)
                    if watch:
                        env.render()
                    reward = reward - 10 if done else reward
                    if info['ale.lives'] < lives:
                        print('Life lost: %d (%d)' % (curr_episode_length, curr_episode_reward))
                        lives = ['ale.lives']
                        reward = reward - 5
                    curr_episode_reward += reward
                    experiences.append((state_converter(state), encoded_actions[action, :], reward, state_converter(next_state)))
                    curr_episode_length += 1

                    # Train network from experiences saved in the replay buffer
                    exps = [experiences[random.randint(0, len(experiences)-1)] for i in range(0, BATCH_SIZE)]
                    sess.run(train_op, feed_dict={
                        s1_tf: np.array([s1 for s1, a1, r1, s2 in exps]),
                        a1_tf: np.array([a1 for s1, a1, r1, s2 in exps]),
                        r1_tf: np.array([r1 for s1, a1, r1, s2 in exps], ndmin=2).T,
                        s2_tf: np.array([s2 for s1, a1, r1, s2 in exps]),
                        training: True
                    })
                    state = next_state

                    # DDQN - Update the target network every N updates
                    n_updates_since_updating_target_net += 1
                    if n_updates_since_updating_target_net >= N_STEPS_BETWEEN_TARGET_UPDATES:
                        n_updates_since_updating_target_net = 0
                        updater(sess)

                lengths[curr_episode] = curr_episode_length
                rewards[curr_episode] = curr_episode_reward
                losses[curr_episode] = sess.run(loss, feed_dict={
                    s1_tf: np.array([s1 for s1, a1, r1, s2 in experiences]),
                    a1_tf: np.array([a1 for s1, a1, r1, s2 in experiences]),
                    r1_tf: np.array([r1 for s1, a1, r1, s2 in experiences], ndmin=2).T,
                    s2_tf: np.array([s2 for s1, a1, r1, s2 in experiences]),
                    training: True
                })
                if curr_episode % N_EPISODES_PER_LOG == 0:
                    # if tf_debug is not None:
                    #     print(sess.run(tf_debug, feed_dict={
                    #         s1_tf: np.array([s1 for s1, a1, r1, s2 in experiences]),
                    #         a1_tf: np.array([a1 for s1, a1, r1, s2 in experiences]),
                    #         r1_tf: np.array([r1 for s1, a1, r1, s2 in experiences], ndmin=2).T,
                    #         s2_tf: np.array([s2 for s1, a1, r1, s2 in experiences]),
                    #         training: True
                    #     }))
                    # _, episode_lengths, episode_rewards = \
                    #     run_n_cartpole_simulations(
                    #         25,
                    #         greedy_sampler(sess, best_action_picker, s1_tf, training),
                    #         env)
                    # test_lengths[int(curr_episode / N_EPISODES_PER_LOG)] = np.mean(np.array(episode_lengths))
                    # saved_progress_path = lr_filename(progress_dir, 'test_lengths_%d' % curr_experiment, learning_rate, 'npy')
                    # if not os.path.exists(os.path.dirname(saved_progress_path)):
                    #     os.makedirs(os.path.dirname(saved_progress_path))
                    # np.save(saved_progress_path, test_lengths)
                    saved_model_path = os.path.join(progress_dir, 'best_model_%d' % np.mean(np.array(rewards[curr_episode])))
                    if not os.path.exists(os.path.dirname(saved_model_path)):
                        os.makedirs(os.path.dirname(saved_model_path))
                    saver.save(sess, saved_model_path)

        np.save(lr_filename(progress_dir, 'lengths_%d' % curr_experiment, learning_rate, 'npy'), lengths)
        np.save(lr_filename(progress_dir, 'losses_%d' % curr_experiment, learning_rate, 'npy'), losses)
        np.save(lr_filename(progress_dir, 'rewards_%d' % curr_experiment, learning_rate, 'npy'), rewards)
        print('Saved! (%f)' % (time.time() - start_time))

<|endoftext|>