<commit_before>from pyVmomi import vim
from tasks import WaitForTasks
from tabulate import tabulate
from misc import sizeof_fmt, humanize_time, esx_get_obj, esx_name, esx_objects

VM_DEFAULT_MEMORY = 1024 # MB
VM_DEFAULT_CPU = 1

###########
# HELPERS #
###########

def vm_get(service, name):
    vms = esx_objects(service, vim.VirtualMachine)
    for v in vms:
        if v.name == name:
            return EsxVirtualMachine(service, v)
    return None

def vm_get_all(service):
    l = []
    vms = esx_objects(service, vim.VirtualMachine)
    for v in vms:
        if v.summary.config.template:
            continue
        vm = EsxVirtualMachine(service, v)
        l.append(vm)
    return l

def vm_guess_folder(vm):
    if vm.parent.name != "vm":
        return vm_guess_folder(vm.parent) + vm.parent.name
    return "/"

def vm_list(s, opt):
    pool = EsxVirtualMachinePool(s)
    vms = pool.list()

    tabs = []
    headers = [ "Name", "Status", "Host", "Folder", "OS", "IP", "CPUs", "Mem (MB)", "NIC", "HDD (GB)", "Uptime" ]

    for v in vms:
        # retrieve infos
        vm = v.info()
        vals = [ vm.name, vm.status, vm.host, vm.folder,
                 vm.os, vm.ip, vm.cpu, vm.mem, vm.nic, vm.hd_size, vm.uptime ]
        tabs.append(vals)
        tabs.sort(reverse=False)

    print tabulate(tabs, headers)

def vm_details(s, opt):
    vm = vm_get(s, opt['<name>'])
    if not vm:
        return

    d = vm.details()
    details = {
        'Name': d.name,
        'Instance UUID': d.instance_uuid[0],
        'Bios UUID': d.bios_uuid,
        'Path to VM': d.path,
        'Guest OS id': d.guest_id[0],
        'Guest OS name': d.guest_name,
        'Host': d.host[0],
        'Last booted timestamp': d.ts
    }

    for name, value in details.items():
        print("  {0:{width}{base}}: {1}".format(name, value, width=25, base='s'))

    print("  Devices:")
    print("  --------")
    for dev in d.devices:
        dev_details = {
            'summary': dev.summary,
            'device type': dev.type
        }

        print("  label: {0}".format(dev.label))
        print("  ------------------")
        for name, value in dev_details.items():
            print("    {0:{width}{base}}: {1}".format(name, value, width=15, base='s'))

            ds = dev.ds
            if ds is None:
                continue

            print("    datastore")
            print("        name: {0}".format(ds.ds_name))
            print("        summary")
            summary = {
                'capacity': sizeof_fmt(ds.ds_capacity),
                'freeSpace': sizeof_fmt(ds.ds_freespace),
                'file system': ds.ds_fs,
                'url': ds.ds_url
            }
            for key, val in summary.items():
                print("            {0}: {1}".format(key, val))
            print("    fileName: {0}".format(ds.filename))
            print("  ------------------")

def vm_spwan(service, vm_name, cluster_name, template_name, memory, cpus, net_name, folder_name):
    content = service.RetrieveContent()
    children = content.rootFolder.childEntity
    for child in children:
        if hasattr(child, 'vmFolder'):
            datacenter = child
        else:
            # some other non-datacenter type object
            continue
        vm_folder = datacenter.vmFolder

    if folder_name != None:
        obj_view = content.viewManager.CreateContainerView(content.rootFolder,[vim.Folder],True)
        folder_list = obj_view.view

        for folder in folder_list:
            if folder.name == folder_name:
                vm_folder = folder

    template_vm = esx_get_obj(content, template_name, vim.VirtualMachine)
    devices = []

    if net_name != None:
        pg_obj = esx_get_obj(content, net_name, vim.dvs.DistributedVirtualPortgroup)
        dvs_port_connection = vim.dvs.PortConnection()
        dvs_port_connection.portgroupKey= pg_obj.key
        dvs_port_connection.switchUuid= pg_obj.config.distributedVirtualSwitch.uuid

        for device in template_vm.config.hardware.device:

            if isinstance(device, vim.vm.device.VirtualEthernetCard):

                nicspec = vim.vm.device.VirtualDeviceSpec()

                nicspec.operation = vim.vm.device.VirtualDeviceSpec.Operation.edit
                nicspec.device = device
                nicspec.device.deviceInfo.label = net_name
                nicspec.device.deviceInfo.summary = net_name
                nicspec.device.backing = vim.vm.device.VirtualEthernetCard.DistributedVirtualPortBackingInfo()
                nicspec.device.backing.port = dvs_port_connection

                devices.append(nicspec)

    cluster = esx_get_obj(content, cluster_name, vim.ClusterComputeResource)
    resource_pool = cluster.resourcePool #

    # vm configuration
    vmconf = vim.vm.ConfigSpec()
    vmconf.numCPUs = cpus
    vmconf.memoryMB = memory
    vmconf.cpuHotAddEnabled = True
    vmconf.memoryHotAddEnabled = True
    if net_name != None:
        vmconf.deviceChange = devices

    relospec = vim.vm.RelocateSpec()
    relospec.pool = resource_pool

    clonespec = vim.vm.CloneSpec()
    clonespec.config = vmconf
    clonespec.location = relospec
    clonespec.powerOn = True

    try:
        clone = template_vm.Clone(folder= vm_folder, name=vm_name, spec=clonespec)
        WaitForTasks(service, [clone])
        print "vm %s successfully created" % vm_name
    except err:
        print err

def vm_create(s, opt):
    vm_name = opt['<name>']
    template = opt['<template>']
    net_name = opt['--network']
    memory = int(opt['--mem'])
    if memory is None:
        memory = VM_DEFAULT_MEMORY
    cpus = int(opt['--cpu'])
    if cpus is None:
        cpus = VM_DEFAULT_CPU
    folder = opt['--folder']

    print 'Trying to clone %s to VM %s' % (template, vm_name)
    if esx_get_obj(s.RetrieveContent(), vm_name, vim.VirtualMachine) != None:
        print 'ERROR: %s already exists' % vm_name
    else:
        vm_spwan(s, vm_name, "Cluster1", template, memory, cpus, net_name, folder)

def vm_delete(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.stop()
        vm.destroy()

def vm_start(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.start()

def vm_stop(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.stop()

def vm_reset(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.reset()

def vm_reboot(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.reboot()

def vm_suspend(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.suspend()

def vm_parser(service, opt):
    if   opt['list']    == True: vm_list(service, opt)
    elif opt['details'] == True: vm_details(service, opt)
    elif opt['create']  == True: vm_create(service, opt)
    elif opt['delete']  == True: vm_delete(service, opt)
    elif opt['start']   == True: vm_start(service, opt)
    elif opt['stop']    == True: vm_stop(service, opt)
    elif opt['reset']   == True: vm_reset(service, opt)
    elif opt['reboot']  == True: vm_reboot(service, opt)
    elif opt['suspend'] == True: vm_suspend(service, opt)

###########
# CLASSES #
###########

class EsxVirtualMachineInfo:
    def __init__(self, vm):
        summary = vm.summary
        config = summary.config
        runtime = summary.runtime
        guest = summary.guest
        storage = summary.storage
        stats = summary.quickStats

        self.name = config.name
        self.status = runtime.powerState
        self.host = esx_name(str(runtime.host))
        self.folder = vm_guess_folder(vm)
        self.os = config.guestFullName
        self.hostname = guest.hostName
        self.ip = guest.ipAddress
        self.cpu = config.numCpu
        self.mem = config.memorySizeMB
        self.nic = config.numEthernetCards
        self.hd_size = (storage.committed + storage.uncommitted) / 1024 / 1024 / 1024
        self.uptime = humanize_time(stats.uptimeSeconds)

class EsxVirtualMachineDeviceHDD:
    def __init__(self, ds, filename):
        self.ds_name = ds.name
        self.ds_capacity = ds.summary.capacity
        self.ds_freespace = ds.summary.freeSpace
        self.ds_fs = ds.summary.type
        self.ds_url = ds.summary.url
        self.filename = filename

class EsxVirtualMachineDevice:
    def __init__(self, d):
        self.summary = d.deviceInfo.summary
        self.type = type(d).__name__
        self.label = d.deviceInfo.label

        self.ds = None
        if d.backing:
            # the following is a bit of a hack, but it lets us build a summary
            # without making many assumptions about the backing type, if the
            # backing type has a file name we *know* it's sitting on a datastore
            # and will have to have all of the following attributes.
            if hasattr(d.backing, 'fileName'):
                datastore = d.backing.datastore
                if datastore:
                    self.ds = EsxVirtualMachineDeviceHDD(datastore, d.backing.fileName)

class EsxVirtualMachineDetails:
    def __init__(self, vm):
        self.name = vm.summary.config.name
        self.instance_uuid = vm.summary.config.instanceUuid,
        self.bios_uuid = vm.summary.config.uuid
        self.path = vm.summary.config.vmPathName
        self.guest_id = vm.summary.config.guestId,
        self.guest_name = vm.summary.config.guestFullName
        self.host = vm.runtime.host.name,
        self.ts = vm.runtime.bootTime

        self.devices = []
        for dev in vm.config.hardware.device:
            d = EsxVirtualMachineDevice(dev)
            self.devices.append(d)

class EsxVirtualMachine:
    def __init__(self, service, vm):
        self.service = service
        self.vm = vm
        self.name = self.vm.name

    def __str__(self):
        return self.name

    def info(self):
        return EsxVirtualMachineInfo(self.vm)

    def details(self):
        return EsxVirtualMachineDetails(self.vm)

    def _task(self, name, t):
        print '{0} VM {1}'.format(name, self.name)
        WaitForTasks(self.service, [t])

    def destroy(self):
        self._task('Destroying', self.vm.Destroy_Task())

    def start(self):
        self._task('Starting', self.vm.PowerOnVM_Task())

    def stop(self):
        self._task('Stopping', self.vm.PowerOffVM_Task())

    def reset(self):
        self._task('Hard Reseting', self.vm.ResetVM_Task())

    def reboot(self):
        print 'Soft Rebooting VM %s' % self.name
        self.vm.RebootGuest()

    def suspend(self):
        self._task('Suspending', self.vm.SuspendVM_Task())

class EsxVirtualMachinePool:
    def __init__(self, service):
        self.vms = vm_get_all(service)

    def list(self):
        return self.vms

    def get(self, name):
        for vm in self.vms:
            if vm.name == name:
                return vm
        return None

    def __str__(self):
        r  = "ESXi Virtual Machines:\n"
        for t in self.vms:
            r += str(t)
        r += "\n"
        return r
<commit_msg>spell check<commit_after>from pyVmomi import vim
from tasks import WaitForTasks
from tabulate import tabulate
from misc import sizeof_fmt, humanize_time, esx_get_obj, esx_name, esx_objects

VM_DEFAULT_MEMORY = 1024 # MB
VM_DEFAULT_CPU = 1

###########
# HELPERS #
###########

def vm_get(service, name):
    vms = esx_objects(service, vim.VirtualMachine)
    for v in vms:
        if v.name == name:
            return EsxVirtualMachine(service, v)
    return None

def vm_get_all(service):
    l = []
    vms = esx_objects(service, vim.VirtualMachine)
    for v in vms:
        if v.summary.config.template:
            continue
        vm = EsxVirtualMachine(service, v)
        l.append(vm)
    return l

def vm_guess_folder(vm):
    if vm.parent.name != "vm":
        return vm_guess_folder(vm.parent) + vm.parent.name
    return "/"

def vm_list(s, opt):
    pool = EsxVirtualMachinePool(s)
    vms = pool.list()

    tabs = []
    headers = [ "Name", "Status", "Host", "Folder", "OS", "IP", "CPUs", "Mem (MB)", "NIC", "HDD (GB)", "Uptime" ]

    for v in vms:
        # retrieve infos
        vm = v.info()
        vals = [ vm.name, vm.status, vm.host, vm.folder,
                 vm.os, vm.ip, vm.cpu, vm.mem, vm.nic, vm.hd_size, vm.uptime ]
        tabs.append(vals)
        tabs.sort(reverse=False)

    print tabulate(tabs, headers)

def vm_details(s, opt):
    vm = vm_get(s, opt['<name>'])
    if not vm:
        return

    d = vm.details()
    details = {
        'Name': d.name,
        'Instance UUID': d.instance_uuid[0],
        'Bios UUID': d.bios_uuid,
        'Path to VM': d.path,
        'Guest OS id': d.guest_id[0],
        'Guest OS name': d.guest_name,
        'Host': d.host[0],
        'Last booted timestamp': d.ts
    }

    for name, value in details.items():
        print("  {0:{width}{base}}: {1}".format(name, value, width=25, base='s'))

    print("  Devices:")
    print("  --------")
    for dev in d.devices:
        dev_details = {
            'summary': dev.summary,
            'device type': dev.type
        }

        print("  label: {0}".format(dev.label))
        print("  ------------------")
        for name, value in dev_details.items():
            print("    {0:{width}{base}}: {1}".format(name, value, width=15, base='s'))

            ds = dev.ds
            if ds is None:
                continue

            print("    datastore")
            print("        name: {0}".format(ds.ds_name))
            print("        summary")
            summary = {
                'capacity': sizeof_fmt(ds.ds_capacity),
                'freeSpace': sizeof_fmt(ds.ds_freespace),
                'file system': ds.ds_fs,
                'url': ds.ds_url
            }
            for key, val in summary.items():
                print("            {0}: {1}".format(key, val))
            print("    fileName: {0}".format(ds.filename))
            print("  ------------------")

def vm_spawn(service, vm_name, cluster_name, template_name, memory, cpus, net_name, folder_name):
    content = service.RetrieveContent()
    children = content.rootFolder.childEntity
    for child in children:
        if hasattr(child, 'vmFolder'):
            datacenter = child
        else:
            # some other non-datacenter type object
            continue
        vm_folder = datacenter.vmFolder

    if folder_name != None:
        obj_view = content.viewManager.CreateContainerView(content.rootFolder,[vim.Folder],True)
        folder_list = obj_view.view

        for folder in folder_list:
            if folder.name == folder_name:
                vm_folder = folder

    template_vm = esx_get_obj(content, template_name, vim.VirtualMachine)
    devices = []

    if net_name != None:
        pg_obj = esx_get_obj(content, net_name, vim.dvs.DistributedVirtualPortgroup)
        dvs_port_connection = vim.dvs.PortConnection()
        dvs_port_connection.portgroupKey= pg_obj.key
        dvs_port_connection.switchUuid= pg_obj.config.distributedVirtualSwitch.uuid

        for device in template_vm.config.hardware.device:

            if isinstance(device, vim.vm.device.VirtualEthernetCard):

                nicspec = vim.vm.device.VirtualDeviceSpec()

                nicspec.operation = vim.vm.device.VirtualDeviceSpec.Operation.edit
                nicspec.device = device
                nicspec.device.deviceInfo.label = net_name
                nicspec.device.deviceInfo.summary = net_name
                nicspec.device.backing = vim.vm.device.VirtualEthernetCard.DistributedVirtualPortBackingInfo()
                nicspec.device.backing.port = dvs_port_connection

                devices.append(nicspec)

    cluster = esx_get_obj(content, cluster_name, vim.ClusterComputeResource)
    resource_pool = cluster.resourcePool #

    # vm configuration
    vmconf = vim.vm.ConfigSpec()
    vmconf.numCPUs = cpus
    vmconf.memoryMB = memory
    vmconf.cpuHotAddEnabled = True
    vmconf.memoryHotAddEnabled = True
    if net_name != None:
        vmconf.deviceChange = devices

    relospec = vim.vm.RelocateSpec()
    relospec.pool = resource_pool

    clonespec = vim.vm.CloneSpec()
    clonespec.config = vmconf
    clonespec.location = relospec
    clonespec.powerOn = True

    try:
        clone = template_vm.Clone(folder= vm_folder, name=vm_name, spec=clonespec)
        WaitForTasks(service, [clone])
        print "vm %s successfully created" % vm_name
    except err:
        print err

def vm_create(s, opt):
    vm_name = opt['<name>']
    template = opt['<template>']
    net_name = opt['--network']
    memory = int(opt['--mem'])
    if memory is None:
        memory = VM_DEFAULT_MEMORY
    cpus = int(opt['--cpu'])
    if cpus is None:
        cpus = VM_DEFAULT_CPU
    folder = opt['--folder']

    print 'Trying to clone %s to VM %s' % (template, vm_name)
    if esx_get_obj(s.RetrieveContent(), vm_name, vim.VirtualMachine) != None:
        print 'ERROR: %s already exists' % vm_name
    else:
        vm_spawn(s, vm_name, "Cluster1", template, memory, cpus, net_name, folder)

def vm_delete(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.stop()
        vm.destroy()

def vm_start(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.start()

def vm_stop(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.stop()

def vm_reset(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.reset()

def vm_reboot(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.reboot()

def vm_suspend(s, opt):
    vm = vm_get(s, opt['<name>'])
    if vm:
        vm.suspend()

def vm_parser(service, opt):
    if   opt['list']    == True: vm_list(service, opt)
    elif opt['details'] == True: vm_details(service, opt)
    elif opt['create']  == True: vm_create(service, opt)
    elif opt['delete']  == True: vm_delete(service, opt)
    elif opt['start']   == True: vm_start(service, opt)
    elif opt['stop']    == True: vm_stop(service, opt)
    elif opt['reset']   == True: vm_reset(service, opt)
    elif opt['reboot']  == True: vm_reboot(service, opt)
    elif opt['suspend'] == True: vm_suspend(service, opt)

###########
# CLASSES #
###########

class EsxVirtualMachineInfo:
    def __init__(self, vm):
        summary = vm.summary
        config = summary.config
        runtime = summary.runtime
        guest = summary.guest
        storage = summary.storage
        stats = summary.quickStats

        self.name = config.name
        self.status = runtime.powerState
        self.host = esx_name(str(runtime.host))
        self.folder = vm_guess_folder(vm)
        self.os = config.guestFullName
        self.hostname = guest.hostName
        self.ip = guest.ipAddress
        self.cpu = config.numCpu
        self.mem = config.memorySizeMB
        self.nic = config.numEthernetCards
        self.hd_size = (storage.committed + storage.uncommitted) / 1024 / 1024 / 1024
        self.uptime = humanize_time(stats.uptimeSeconds)

class EsxVirtualMachineDeviceHDD:
    def __init__(self, ds, filename):
        self.ds_name = ds.name
        self.ds_capacity = ds.summary.capacity
        self.ds_freespace = ds.summary.freeSpace
        self.ds_fs = ds.summary.type
        self.ds_url = ds.summary.url
        self.filename = filename

class EsxVirtualMachineDevice:
    def __init__(self, d):
        self.summary = d.deviceInfo.summary
        self.type = type(d).__name__
        self.label = d.deviceInfo.label

        self.ds = None
        if d.backing:
            # the following is a bit of a hack, but it lets us build a summary
            # without making many assumptions about the backing type, if the
            # backing type has a file name we *know* it's sitting on a datastore
            # and will have to have all of the following attributes.
            if hasattr(d.backing, 'fileName'):
                datastore = d.backing.datastore
                if datastore:
                    self.ds = EsxVirtualMachineDeviceHDD(datastore, d.backing.fileName)

class EsxVirtualMachineDetails:
    def __init__(self, vm):
        self.name = vm.summary.config.name
        self.instance_uuid = vm.summary.config.instanceUuid,
        self.bios_uuid = vm.summary.config.uuid
        self.path = vm.summary.config.vmPathName
        self.guest_id = vm.summary.config.guestId,
        self.guest_name = vm.summary.config.guestFullName
        self.host = vm.runtime.host.name,
        self.ts = vm.runtime.bootTime

        self.devices = []
        for dev in vm.config.hardware.device:
            d = EsxVirtualMachineDevice(dev)
            self.devices.append(d)

class EsxVirtualMachine:
    def __init__(self, service, vm):
        self.service = service
        self.vm = vm
        self.name = self.vm.name

    def __str__(self):
        return self.name

    def info(self):
        return EsxVirtualMachineInfo(self.vm)

    def details(self):
        return EsxVirtualMachineDetails(self.vm)

    def _task(self, name, t):
        print '{0} VM {1}'.format(name, self.name)
        WaitForTasks(self.service, [t])

    def destroy(self):
        self._task('Destroying', self.vm.Destroy_Task())

    def start(self):
        self._task('Starting', self.vm.PowerOnVM_Task())

    def stop(self):
        self._task('Stopping', self.vm.PowerOffVM_Task())

    def reset(self):
        self._task('Hard Reseting', self.vm.ResetVM_Task())

    def reboot(self):
        print 'Soft Rebooting VM %s' % self.name
        self.vm.RebootGuest()

    def suspend(self):
        self._task('Suspending', self.vm.SuspendVM_Task())

class EsxVirtualMachinePool:
    def __init__(self, service):
        self.vms = vm_get_all(service)

    def list(self):
        return self.vms

    def get(self, name):
        for vm in self.vms:
            if vm.name == name:
                return vm
        return None

    def __str__(self):
        r  = "ESXi Virtual Machines:\n"
        for t in self.vms:
            r += str(t)
        r += "\n"
        return r
<|endoftext|>


<commit_before>## \namespace orbit::sc1d
## \brief The classes and functions for foils
##
## Classes:
##  - sc1DNode -longitudinal space charge node for the TEAPOT lattices
##
## addLongitudinalSpaceChargeNode- function to add one longitudinal space charge node to the lattice

from orbit.space_charge.sc1d.sc1DNode import SC1D_AccNode
from orbit.space_charge.sc1d.sc1DNode import FreqDep_SC1D_AccNode
from orbit.space_charge.sc1d.sc1DNode import BetFreqDep_SC1D_AccNode

from orbit.space_charge.sc1d.scLatticeModifications import addLongitudinalSpaceChargeNode

__all__ = []
__all__.append("sc1DNode")
__all__.append("addLongitudinalSpaceChargeNode")
<commit_msg>Added a routine to make the longitudinal space charge operation a child node of another node.<commit_after>## \namespace orbit::sc1d
## \brief The classes and functions for foils
##
## Classes:
##  - sc1DNode -longitudinal space charge node for the TEAPOT lattices
##
## addLongitudinalSpaceChargeNode- function to add one longitudinal space charge node to the lattice

from orbit.space_charge.sc1d.sc1DNode import SC1D_AccNode
from orbit.space_charge.sc1d.sc1DNode import FreqDep_SC1D_AccNode
from orbit.space_charge.sc1d.sc1DNode import BetFreqDep_SC1D_AccNode

from orbit.space_charge.sc1d.scLatticeModifications import addLongitudinalSpaceChargeNode
from orbit.space_charge.sc1d.scLatticeModifications import addLongitudinalSpaceChargeNodeAsChild

__all__ = []
__all__.append("sc1DNode")
__all__.append("addLongitudinalSpaceChargeNode")
__all__.append("addLongitudinalSpaceChargeNodeAsChild")
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2016-2020 CNRS

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr

"""
#########
Protocols
#########

"""

import warnings
import collections
import threading
import itertools
from typing import Union, Dict, Iterator, Callable, Any, Text, Optional

try:
    from typing import Literal
except ImportError:
    from typing_extensions import Literal

Subset = Literal["train", "development", "test"]
LEGACY_SUBSET_MAPPING = {"train": "trn", "development": "dev", "test": "tst"}

Preprocessor = Callable[["ProtocolFile"], Any]
Preprocessors = Dict[Text, Preprocessor]

class ProtocolFile(collections.abc.MutableMapping):
    """Protocol file with lazy preprocessors

    This is a dict-like data structure where some values may depend on other
    values, and are only computed if/when requested. Once computed, they are
    cached and never recomputed again.

    Parameters
    ----------
    precomputed : dict
        Regular dictionary with precomputed values
    lazy : dict, optional
        Dictionary describing how lazy value needs to be computed.
        Values are callable expecting a dictionary as input and returning the
        computed value.

    """

    def __init__(self, precomputed: Union[Dict, "ProtocolFile"], lazy: Dict = None):

        if lazy is None:
            lazy = dict()

        if isinstance(precomputed, ProtocolFile):
            # when 'precomputed' is a ProtocolFile, it may already contain lazy keys.

            # we use 'precomputed' precomputed keys as precomputed keys
            self._store: Dict = abs(precomputed)

            # we handle the corner case where the intersection of 'precomputed' lazy keys
            # and 'lazy' keys is not empty. this is currently achieved by "unlazying" the
            # 'precomputed' one (which is probably not the most efficient solution).
            for key in set(precomputed.lazy) & set(lazy):
                self._store[key] = precomputed[key]

            # we use the union of 'precomputed' lazy keys and provided 'lazy' keys as lazy keys
            compound_lazy = dict(precomputed.lazy)
            compound_lazy.update(lazy)
            self.lazy: Dict = compound_lazy

        else:
            # when 'precomputed' is a Dict, we use it directly as precomputed keys
            # and 'lazy' as lazy keys.
            self._store = dict(precomputed)
            self.lazy = dict(lazy)

        # re-entrant lock used below to make ProtocolFile thread-safe
        self.lock_ = threading.RLock()

        # this is needed to avoid infinite recursion
        # when a key is both in precomputed and lazy.
        # keys with evaluating_ > 0 are currently being evaluated
        # and therefore should be taken from precomputed
        self.evaluating_ = collections.Counter()

    def __abs__(self):
        with self.lock_:
            return dict(self._store)

    def __getitem__(self, key):
        with self.lock_:

            if key in self.lazy and self.evaluating_[key] == 0:

                # mark lazy key as being evaluated
                self.evaluating_.update([key])

                # apply preprocessor once and remove it
                value = self.lazy[key](self)
                del self.lazy[key]

                # warn the user when a precomputed key is modified
                if key in self._store and value != self._store[key]:
                    msg = 'Existing precomputed key "{key}" has been modified by a preprocessor.'
                    warnings.warn(msg.format(key=key))

                # store the output of the lazy computation
                # so that it is available for future access
                self._store[key] = value

                # lazy evaluation is finished for key
                self.evaluating_.subtract([key])

            return self._store[key]

    def __setitem__(self, key, value):
        with self.lock_:

            if key in self.lazy:
                del self.lazy[key]

            self._store[key] = value

    def __delitem__(self, key):
        with self.lock_:

            if key in self.lazy:
                del self.lazy[key]

            del self._store[key]

    def __iter__(self):
        with self.lock_:

            for key in self._store:
                yield key

            lazy_keys = list(self.lazy)
            for key in lazy_keys:
                if key in self._store:
                    continue
                yield key

    def __len__(self):
        with self.lock_:
            return len(set(self._store) | set(self.lazy))

    def files(self) -> Iterator["ProtocolFile"]:
        """Iterate over all files

        When `current_file` refers to only one file,
            yield it and return.
        When `current_file` refers to a list of file (i.e. 'uri' is a list),
            yield each file separately.

        Examples
        --------
        >>> current_file = ProtocolFile({
        ...     'uri': 'my_uri',
        ...     'database': 'my_database'})
        >>> for file in current_file.files():
        ...     print(file['uri'], file['database'])
        my_uri my_database

        >>> current_file = {
        ...     'uri': ['my_uri1', 'my_uri2', 'my_uri3'],
        ...     'database': 'my_database'}
        >>> for file in current_file.files():
        ...     print(file['uri'], file['database'])
        my_uri1 my_database
        my_uri2 my_database
        my_uri3 my_database

        """

        uris = self["uri"]
        if not isinstance(uris, list):
            yield self
            return

        n_uris = len(uris)

        # iterate over precomputed keys and make sure

        precomputed = {"uri": uris}
        for key, value in abs(self).items():

            if key == "uri":
                continue

            if not isinstance(value, list):
                precomputed[key] = itertools.repeat(value)

            else:
                if len(value) != n_uris:
                    msg = (
                        f'Mismatch between number of "uris" ({n_uris}) '
                        f'and number of "{key}" ({len(value)}).'
                    )
                    raise ValueError(msg)
                precomputed[key] = value

        keys = list(precomputed.keys())
        for values in zip(*precomputed.values()):
            precomputed_one = dict(zip(keys, values))
            yield ProtocolFile(precomputed_one, self.lazy)

class Protocol:
    """Experimental protocol

    An experimental protocol usually defines three subsets: a training subset,
    a development subset, and a test subset.

    An experimental protocol can be defined programmatically by creating a
    class that inherits from SpeakerDiarizationProtocol and implements at least
    one of `train_iter`, `development_iter` and `test_iter` methods:

        >>> class MyProtocol(Protocol):
        ...     def train_iter(self) -> Iterator[Dict]:
        ...         yield {"uri": "filename1", "any_other_key": "..."}
        ...         yield {"uri": "filename2", "any_other_key": "..."}

    `{subset}_iter` should return an iterator of dictionnaries with
        - "uri" key (mandatory) that provides a unique file identifier (usually
          the filename),
        - any other key that the protocol may provide.

    It can then be used in Python like this:

        >>> protocol = MyProtocol()
        >>> for file in protocol.train():
        ...    print(file["uri"])
        filename1
        filename2

    An experimental protocol can also be defined using `pyannote.database`
    configuration file, whose (configurable) path defaults to "~/database.yml".

    ~~~ Content of ~/database.yml ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    Protocols:
      MyDatabase:
        Protocol:
          MyProtocol:
            train:
                uri: /path/to/collection.lst
                any_other_key: ... # see custom loader documentation
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    where "/path/to/collection.lst" contains the list of identifiers of the
    files in the collection:

    ~~~ Content of "/path/to/collection.lst ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    filename1
    filename2
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    It can then be used in Python like this:

        >>> from pyannote.database import get_protocol
        >>> protocol = get_protocol('MyDatabase.Protocol.MyProtocol')
        >>> for file in protocol.train():
        ...    print(file["uri"])
        filename1
        filename2

    This class is usually inherited from, but can be used directly.

    Parameters
    ----------
    preprocessors : dict
        Preprocess protocol files so that `file[key] = preprocessors[key](file)`
        for each key in `preprocessors`. In case `preprocessors[key]` is not
        callable, it should be a string containing placeholders for `file` keys
        (e.g. {'audio': '/path/to/{uri}.wav'})
    """

    def __init__(self, preprocessors: Optional[Preprocessors] = None):
        super().__init__()

        if preprocessors is None:
            preprocessors = dict()

        self.preprocessors = dict()
        for key, preprocessor in preprocessors.items():

            if callable(preprocessor):
                self.preprocessors[key] = preprocessor

            # when `preprocessor` is not callable, it should be a string
            # containing placeholder for item key (e.g. '/path/to/{uri}.wav')
            elif isinstance(preprocessor, str):
                preprocessor_copy = str(preprocessor)

                def func(current_file):
                    return preprocessor_copy.format(**current_file)

                self.preprocessors[key] = func

            else:
                msg = f'"{key}" preprocessor is neither a callable nor a string.'
                raise ValueError(msg)

    def preprocess(self, current_file: Union[Dict, ProtocolFile]) -> ProtocolFile:
        return ProtocolFile(current_file, lazy=self.preprocessors)

    def __str__(self):
        return self.__doc__

    def train_iter(self) -> Iterator[Union[Dict, ProtocolFile]]:
        """Iterate over files in the training subset"""
        raise NotImplementedError()

    def development_iter(self) -> Iterator[Union[Dict, ProtocolFile]]:
        """Iterate over files in the development subset"""
        raise NotImplementedError()

    def test_iter(self) -> Iterator[Union[Dict, ProtocolFile]]:
        """Iterate over files in the test subset"""
        raise NotImplementedError()

    def subset_helper(self, subset: Subset) -> Iterator[ProtocolFile]:

        try:
            files = getattr(self, f"{subset}_iter")()
        except (AttributeError, NotImplementedError):
            # previous pyannote.database versions used `trn_iter` instead of
            # `train_iter`, `dev_iter` instead of `development_iter`, and
            # `tst_iter` instead of `test_iter`. therefore, we use the legacy
            # version when it is available (and the new one is not).
            subset_legacy = LEGACY_SUBSET_MAPPING[subset]
            try:
                files = getattr(self, f"{subset_legacy}_iter")()
            except AttributeError:
                msg = f"Protocol does not implement a {subset} subset."
                raise NotImplementedError(msg)

        for file in files:
            yield self.preprocess(file)

    def train(self) -> Iterator[ProtocolFile]:
        return self.subset_helper("train")

    def development(self) -> Iterator[ProtocolFile]:
        return self.subset_helper("development")

    def test(self) -> Iterator[ProtocolFile]:
        return self.subset_helper("test")

    def files(self) -> Iterator[ProtocolFile]:
        """Iterate over all files in `protocol`"""

        # imported here to avoid circular imports
        from pyannote.database.util import get_unique_identifier

        yielded_uris = set()

        for method in [
            "development",
            "development_enrolment",
            "development_trial",
            "test",
            "test_enrolment",
            "test_trial",
            "train",
            "train_enrolment",
            "train_trial",
        ]:

            if not hasattr(self, method):
                continue

            def iterate():
                try:
                    for file in getattr(self, method)():
                        yield file
                except (AttributeError, NotImplementedError):
                    return

            for current_file in iterate():

                # skip "files" that do not contain a "uri" entry.
                # this happens for speaker verification trials that contain
                # two nested files "file1" and "file2"
                # see https://github.com/pyannote/pyannote-db-voxceleb/issues/4
                if "uri" not in current_file:
                    continue

                for current_file_ in current_file.files():

                    # corner case when the same file is yielded several times
                    uri = get_unique_identifier(current_file_)
                    if uri in yielded_uris:
                        continue

                    yield current_file_

                    yielded_uris.add(uri)
<commit_msg>fix: fix ProtocolFile.__iter__ (dictionary changed size during iteration)<commit_after>#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2016-2020 CNRS

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr

"""
#########
Protocols
#########

"""

import warnings
import collections
import threading
import itertools
from typing import Union, Dict, Iterator, Callable, Any, Text, Optional

try:
    from typing import Literal
except ImportError:
    from typing_extensions import Literal

Subset = Literal["train", "development", "test"]
LEGACY_SUBSET_MAPPING = {"train": "trn", "development": "dev", "test": "tst"}

Preprocessor = Callable[["ProtocolFile"], Any]
Preprocessors = Dict[Text, Preprocessor]

class ProtocolFile(collections.abc.MutableMapping):
    """Protocol file with lazy preprocessors

    This is a dict-like data structure where some values may depend on other
    values, and are only computed if/when requested. Once computed, they are
    cached and never recomputed again.

    Parameters
    ----------
    precomputed : dict
        Regular dictionary with precomputed values
    lazy : dict, optional
        Dictionary describing how lazy value needs to be computed.
        Values are callable expecting a dictionary as input and returning the
        computed value.

    """

    def __init__(self, precomputed: Union[Dict, "ProtocolFile"], lazy: Dict = None):

        if lazy is None:
            lazy = dict()

        if isinstance(precomputed, ProtocolFile):
            # when 'precomputed' is a ProtocolFile, it may already contain lazy keys.

            # we use 'precomputed' precomputed keys as precomputed keys
            self._store: Dict = abs(precomputed)

            # we handle the corner case where the intersection of 'precomputed' lazy keys
            # and 'lazy' keys is not empty. this is currently achieved by "unlazying" the
            # 'precomputed' one (which is probably not the most efficient solution).
            for key in set(precomputed.lazy) & set(lazy):
                self._store[key] = precomputed[key]

            # we use the union of 'precomputed' lazy keys and provided 'lazy' keys as lazy keys
            compound_lazy = dict(precomputed.lazy)
            compound_lazy.update(lazy)
            self.lazy: Dict = compound_lazy

        else:
            # when 'precomputed' is a Dict, we use it directly as precomputed keys
            # and 'lazy' as lazy keys.
            self._store = dict(precomputed)
            self.lazy = dict(lazy)

        # re-entrant lock used below to make ProtocolFile thread-safe
        self.lock_ = threading.RLock()

        # this is needed to avoid infinite recursion
        # when a key is both in precomputed and lazy.
        # keys with evaluating_ > 0 are currently being evaluated
        # and therefore should be taken from precomputed
        self.evaluating_ = collections.Counter()

    def __abs__(self):
        with self.lock_:
            return dict(self._store)

    def __getitem__(self, key):
        with self.lock_:

            if key in self.lazy and self.evaluating_[key] == 0:

                # mark lazy key as being evaluated
                self.evaluating_.update([key])

                # apply preprocessor once and remove it
                value = self.lazy[key](self)
                del self.lazy[key]

                # warn the user when a precomputed key is modified
                if key in self._store and value != self._store[key]:
                    msg = 'Existing precomputed key "{key}" has been modified by a preprocessor.'
                    warnings.warn(msg.format(key=key))

                # store the output of the lazy computation
                # so that it is available for future access
                self._store[key] = value

                # lazy evaluation is finished for key
                self.evaluating_.subtract([key])

            return self._store[key]

    def __setitem__(self, key, value):
        with self.lock_:

            if key in self.lazy:
                del self.lazy[key]

            self._store[key] = value

    def __delitem__(self, key):
        with self.lock_:

            if key in self.lazy:
                del self.lazy[key]

            del self._store[key]

    def __iter__(self):
        with self.lock_:

            store_keys = list(self._store)
            for key in store_keys:
                yield key

            lazy_keys = list(self.lazy)
            for key in lazy_keys:
                if key in self._store:
                    continue
                yield key

    def __len__(self):
        with self.lock_:
            return len(set(self._store) | set(self.lazy))

    def files(self) -> Iterator["ProtocolFile"]:
        """Iterate over all files

        When `current_file` refers to only one file,
            yield it and return.
        When `current_file` refers to a list of file (i.e. 'uri' is a list),
            yield each file separately.

        Examples
        --------
        >>> current_file = ProtocolFile({
        ...     'uri': 'my_uri',
        ...     'database': 'my_database'})
        >>> for file in current_file.files():
        ...     print(file['uri'], file['database'])
        my_uri my_database

        >>> current_file = {
        ...     'uri': ['my_uri1', 'my_uri2', 'my_uri3'],
        ...     'database': 'my_database'}
        >>> for file in current_file.files():
        ...     print(file['uri'], file['database'])
        my_uri1 my_database
        my_uri2 my_database
        my_uri3 my_database

        """

        uris = self["uri"]
        if not isinstance(uris, list):
            yield self
            return

        n_uris = len(uris)

        # iterate over precomputed keys and make sure

        precomputed = {"uri": uris}
        for key, value in abs(self).items():

            if key == "uri":
                continue

            if not isinstance(value, list):
                precomputed[key] = itertools.repeat(value)

            else:
                if len(value) != n_uris:
                    msg = (
                        f'Mismatch between number of "uris" ({n_uris}) '
                        f'and number of "{key}" ({len(value)}).'
                    )
                    raise ValueError(msg)
                precomputed[key] = value

        keys = list(precomputed.keys())
        for values in zip(*precomputed.values()):
            precomputed_one = dict(zip(keys, values))
            yield ProtocolFile(precomputed_one, self.lazy)

class Protocol:
    """Experimental protocol

    An experimental protocol usually defines three subsets: a training subset,
    a development subset, and a test subset.

    An experimental protocol can be defined programmatically by creating a
    class that inherits from SpeakerDiarizationProtocol and implements at least
    one of `train_iter`, `development_iter` and `test_iter` methods:

        >>> class MyProtocol(Protocol):
        ...     def train_iter(self) -> Iterator[Dict]:
        ...         yield {"uri": "filename1", "any_other_key": "..."}
        ...         yield {"uri": "filename2", "any_other_key": "..."}

    `{subset}_iter` should return an iterator of dictionnaries with
        - "uri" key (mandatory) that provides a unique file identifier (usually
          the filename),
        - any other key that the protocol may provide.

    It can then be used in Python like this:

        >>> protocol = MyProtocol()
        >>> for file in protocol.train():
        ...    print(file["uri"])
        filename1
        filename2

    An experimental protocol can also be defined using `pyannote.database`
    configuration file, whose (configurable) path defaults to "~/database.yml".

    ~~~ Content of ~/database.yml ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    Protocols:
      MyDatabase:
        Protocol:
          MyProtocol:
            train:
                uri: /path/to/collection.lst
                any_other_key: ... # see custom loader documentation
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    where "/path/to/collection.lst" contains the list of identifiers of the
    files in the collection:

    ~~~ Content of "/path/to/collection.lst ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    filename1
    filename2
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    It can then be used in Python like this:

        >>> from pyannote.database import get_protocol
        >>> protocol = get_protocol('MyDatabase.Protocol.MyProtocol')
        >>> for file in protocol.train():
        ...    print(file["uri"])
        filename1
        filename2

    This class is usually inherited from, but can be used directly.

    Parameters
    ----------
    preprocessors : dict
        Preprocess protocol files so that `file[key] = preprocessors[key](file)`
        for each key in `preprocessors`. In case `preprocessors[key]` is not
        callable, it should be a string containing placeholders for `file` keys
        (e.g. {'audio': '/path/to/{uri}.wav'})
    """

    def __init__(self, preprocessors: Optional[Preprocessors] = None):
        super().__init__()

        if preprocessors is None:
            preprocessors = dict()

        self.preprocessors = dict()
        for key, preprocessor in preprocessors.items():

            if callable(preprocessor):
                self.preprocessors[key] = preprocessor

            # when `preprocessor` is not callable, it should be a string
            # containing placeholder for item key (e.g. '/path/to/{uri}.wav')
            elif isinstance(preprocessor, str):
                preprocessor_copy = str(preprocessor)

                def func(current_file):
                    return preprocessor_copy.format(**current_file)

                self.preprocessors[key] = func

            else:
                msg = f'"{key}" preprocessor is neither a callable nor a string.'
                raise ValueError(msg)

    def preprocess(self, current_file: Union[Dict, ProtocolFile]) -> ProtocolFile:
        return ProtocolFile(current_file, lazy=self.preprocessors)

    def __str__(self):
        return self.__doc__

    def train_iter(self) -> Iterator[Union[Dict, ProtocolFile]]:
        """Iterate over files in the training subset"""
        raise NotImplementedError()

    def development_iter(self) -> Iterator[Union[Dict, ProtocolFile]]:
        """Iterate over files in the development subset"""
        raise NotImplementedError()

    def test_iter(self) -> Iterator[Union[Dict, ProtocolFile]]:
        """Iterate over files in the test subset"""
        raise NotImplementedError()

    def subset_helper(self, subset: Subset) -> Iterator[ProtocolFile]:

        try:
            files = getattr(self, f"{subset}_iter")()
        except (AttributeError, NotImplementedError):
            # previous pyannote.database versions used `trn_iter` instead of
            # `train_iter`, `dev_iter` instead of `development_iter`, and
            # `tst_iter` instead of `test_iter`. therefore, we use the legacy
            # version when it is available (and the new one is not).
            subset_legacy = LEGACY_SUBSET_MAPPING[subset]
            try:
                files = getattr(self, f"{subset_legacy}_iter")()
            except AttributeError:
                msg = f"Protocol does not implement a {subset} subset."
                raise NotImplementedError(msg)

        for file in files:
            yield self.preprocess(file)

    def train(self) -> Iterator[ProtocolFile]:
        return self.subset_helper("train")

    def development(self) -> Iterator[ProtocolFile]:
        return self.subset_helper("development")

    def test(self) -> Iterator[ProtocolFile]:
        return self.subset_helper("test")

    def files(self) -> Iterator[ProtocolFile]:
        """Iterate over all files in `protocol`"""

        # imported here to avoid circular imports
        from pyannote.database.util import get_unique_identifier

        yielded_uris = set()

        for method in [
            "development",
            "development_enrolment",
            "development_trial",
            "test",
            "test_enrolment",
            "test_trial",
            "train",
            "train_enrolment",
            "train_trial",
        ]:

            if not hasattr(self, method):
                continue

            def iterate():
                try:
                    for file in getattr(self, method)():
                        yield file
                except (AttributeError, NotImplementedError):
                    return

            for current_file in iterate():

                # skip "files" that do not contain a "uri" entry.
                # this happens for speaker verification trials that contain
                # two nested files "file1" and "file2"
                # see https://github.com/pyannote/pyannote-db-voxceleb/issues/4
                if "uri" not in current_file:
                    continue

                for current_file_ in current_file.files():

                    # corner case when the same file is yielded several times
                    uri = get_unique_identifier(current_file_)
                    if uri in yielded_uris:
                        continue

                    yield current_file_

                    yielded_uris.add(uri)
<|endoftext|>


<commit_before>'''
Created on 06.12.2013

@author: jan-hendrikprinz

This module provides the unittest for the pcca module

'''
import unittest
import numpy as np
from scipy import sparse

from pyemma.util.numeric import assert_allclose
from pcca import pcca

class TestPCCA(unittest.TestCase):

    def setUp(self):
        pass

    def test_pcca_no_transition_matrix(self):
        P = np.array([[1.0, 1.0],
                      [0.1, 0.9]])
        try:
            pcca(P, 2)
            # no ValueError? then fail.
            assert False
        except ValueError:
            pass
        except: # different exception
            assert False

    def test_pcca_no_detailed_balance(self):
        P = np.array([[0.8, 0.1, 0.1],
                      [0.3, 0.2, 0.5],
                      [0.6, 0.3, 0.1]])
        try:
            pcca(P, 2)
            # no ValueError? then fail.
            assert False
        except ValueError:
            pass
        except: # different exception
            assert False

    def test_pcca_1(self):
        P = np.array([[1, 0],
                      [0, 1]])
        chi = pcca(P, 2)
        sol = np.array([[ 1.,  0.],
                        [ 0.,  1.]])
        assert_allclose(chi, sol)

    def test_pcca_2(self):
        P = np.array([[0.0, 1.0, 0.0],
                      [0.0, 0.999, 0.001],
                      [0.0, 0.001, 0.999]])
        chi = pcca(P, 2)
        sol = np.array([[ 1.,  0.],
                        [ 1.,  0.],
                        [ 0.,  1.]])
        assert_allclose(chi, sol)

    def test_pcca_3(self):
        P = np.array([[0.9, 0.1, 0.0, 0.0],
                      [0.1, 0.9, 0.0, 0.0],
                      [0.0, 0.0, 0.8, 0.2],
                      [0.0, 0.0, 0.2, 0.8]])
        # n=2
        chi = pcca(P, 2)
        sol = np.array([[ 1.,  0.],
                        [ 1.,  0.],
                        [ 0.,  1.],
                        [ 0.,  1.]])
        assert_allclose(chi, sol)
        # n=3
        chi = pcca(P, 3)
        sol = np.array([[ 1.,  0.,  0.],
                        [ 0.,  1.,  0.],
                        [ 0.,  0.,  1.],
                        [ 0.,  0.,  1.]])
        assert_allclose(chi, sol)
        # n=4
        chi = pcca(P, 4)
        sol = np.array([[ 1.,  0.,  0.,  0.],
                        [ 0.,  1.,  0.,  0.],
                        [ 0.,  0.,  1.,  0.],
                        [ 0.,  0.,  0.,  1.]])
        assert_allclose(chi, sol)

    def test_pcca_4(self):
        P = np.array([[0.9, 0.1, 0.0, 0.0],
                      [0.1, 0.8, 0.1, 0.0],
                      [0.0, 0.0, 0.8, 0.2],
                      [0.0, 0.0, 0.2, 0.8]])
        chi = pcca(P, 2)
        sol = np.array([[ 1.,  0.],
                        [ 1.,  0.],
                        [ 1.,  0.],
                        [ 0.,  1.]])
        assert_allclose(chi, sol)

    def test_pcca_5(self):
        P = np.array([[0.9, 0.1, 0.0, 0.0, 0.0],
                      [0.1, 0.9, 0.0, 0.0, 0.0],
                      [0.0, 0.1, 0.8, 0.1, 0.0],
                      [0.0, 0.0, 0.0, 0.8, 0.2],
                      [0.0, 0.0, 0.0, 0.2, 0.8]])
        # n=2
        chi = pcca(P, 2)
        sol = np.array([[ 1. ,  0. ],
                        [ 1. ,  0. ],
                        [ 0.5,  0.5],
                        [ 0. ,  1. ],
                        [ 0. ,  1. ]])
        assert_allclose(chi, sol)
        # n=4
        chi = pcca(P, 4)
        sol = np.array([[ 1. ,  0. ,  0. ,  0. ],
                        [ 0. ,  1. ,  0. ,  0. ],
                        [ 0. ,  0.5,  0.5,  0. ],
                        [ 0. ,  0. ,  1. ,  0. ],
                        [ 0. ,  0. ,  0. ,  1. ]])
        assert_allclose(chi, sol)

    def test_pcca_large(self):
        import os
        P = np.loadtxt(os.path.split(__file__)[0]+'/../tests/P_rev_251x251.dat')
        # n=2
        chi = pcca(P, 2)
        assert(np.alltrue(chi >= 0))
        assert(np.alltrue(chi <= 1))
        # n=3
        chi = pcca(P, 3)
        assert(np.alltrue(chi >= 0))
        assert(np.alltrue(chi <= 1))
        # n=4
        chi = pcca(P, 4)
        assert(np.alltrue(chi >= 0))
        assert(np.alltrue(chi <= 1))

if __name__ == "__main__":
    unittest.main()
<commit_msg>[pcca-test] use assertRaises<commit_after>'''
Created on 06.12.2013

@author: jan-hendrikprinz

This module provides the unittest for the pcca module

'''
import unittest
import numpy as np
from scipy import sparse

from pyemma.util.numeric import assert_allclose
from pcca import pcca

class TestPCCA(unittest.TestCase):

    def test_pcca_no_transition_matrix(self):
        P = np.array([[1.0, 1.0],
                      [0.1, 0.9]])

        with self.assertRaises(ValueError):
            pcca(P, 2)

    def test_pcca_no_detailed_balance(self):
        P = np.array([[0.8, 0.1, 0.1],
                      [0.3, 0.2, 0.5],
                      [0.6, 0.3, 0.1]])
        with self.assertRaises(ValueError):
            pcca(P, 2)

    def test_pcca_1(self):
        P = np.array([[1, 0],
                      [0, 1]])
        chi = pcca(P, 2)
        sol = np.array([[ 1.,  0.],
                        [ 0.,  1.]])
        assert_allclose(chi, sol)

    def test_pcca_2(self):
        P = np.array([[0.0, 1.0, 0.0],
                      [0.0, 0.999, 0.001],
                      [0.0, 0.001, 0.999]])
        chi = pcca(P, 2)
        sol = np.array([[ 1.,  0.],
                        [ 1.,  0.],
                        [ 0.,  1.]])
        assert_allclose(chi, sol)

    def test_pcca_3(self):
        P = np.array([[0.9, 0.1, 0.0, 0.0],
                      [0.1, 0.9, 0.0, 0.0],
                      [0.0, 0.0, 0.8, 0.2],
                      [0.0, 0.0, 0.2, 0.8]])
        # n=2
        chi = pcca(P, 2)
        sol = np.array([[ 1.,  0.],
                        [ 1.,  0.],
                        [ 0.,  1.],
                        [ 0.,  1.]])
        assert_allclose(chi, sol)
        # n=3
        chi = pcca(P, 3)
        sol = np.array([[ 1.,  0.,  0.],
                        [ 0.,  1.,  0.],
                        [ 0.,  0.,  1.],
                        [ 0.,  0.,  1.]])
        assert_allclose(chi, sol)
        # n=4
        chi = pcca(P, 4)
        sol = np.array([[ 1.,  0.,  0.,  0.],
                        [ 0.,  1.,  0.,  0.],
                        [ 0.,  0.,  1.,  0.],
                        [ 0.,  0.,  0.,  1.]])
        assert_allclose(chi, sol)

    def test_pcca_4(self):
        P = np.array([[0.9, 0.1, 0.0, 0.0],
                      [0.1, 0.8, 0.1, 0.0],
                      [0.0, 0.0, 0.8, 0.2],
                      [0.0, 0.0, 0.2, 0.8]])
        chi = pcca(P, 2)
        sol = np.array([[ 1.,  0.],
                        [ 1.,  0.],
                        [ 1.,  0.],
                        [ 0.,  1.]])
        assert_allclose(chi, sol)

    def test_pcca_5(self):
        P = np.array([[0.9, 0.1, 0.0, 0.0, 0.0],
                      [0.1, 0.9, 0.0, 0.0, 0.0],
                      [0.0, 0.1, 0.8, 0.1, 0.0],
                      [0.0, 0.0, 0.0, 0.8, 0.2],
                      [0.0, 0.0, 0.0, 0.2, 0.8]])
        # n=2
        chi = pcca(P, 2)
        sol = np.array([[ 1. ,  0. ],
                        [ 1. ,  0. ],
                        [ 0.5,  0.5],
                        [ 0. ,  1. ],
                        [ 0. ,  1. ]])
        assert_allclose(chi, sol)
        # n=4
        chi = pcca(P, 4)
        sol = np.array([[ 1. ,  0. ,  0. ,  0. ],
                        [ 0. ,  1. ,  0. ,  0. ],
                        [ 0. ,  0.5,  0.5,  0. ],
                        [ 0. ,  0. ,  1. ,  0. ],
                        [ 0. ,  0. ,  0. ,  1. ]])
        assert_allclose(chi, sol)

    def test_pcca_large(self):
        import os
        P = np.loadtxt(os.path.split(__file__)[0]+'/../tests/P_rev_251x251.dat')
        # n=2
        chi = pcca(P, 2)
        assert(np.alltrue(chi >= 0))
        assert(np.alltrue(chi <= 1))
        # n=3
        chi = pcca(P, 3)
        assert(np.alltrue(chi >= 0))
        assert(np.alltrue(chi <= 1))
        # n=4
        chi = pcca(P, 4)
        assert(np.alltrue(chi >= 0))
        assert(np.alltrue(chi <= 1))

if __name__ == "__main__":
    unittest.main()
<|endoftext|>


<commit_before># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

import re
import csv
import collections
import itertools
from io import open
import math
from six.moves import zip
import logging

from monty.json import MSONable, MontyDecoder

import numpy as np

from pymatgen.core.structure import Molecule
from pymatgen.io.qchem.outputs import edges_from_babel, build_MoleculeGraph, is_isomorphic
import networkx as nx

# have_babel = True
# try:
#     from pymatgen.io.babel import BabelMolAdaptor
#     import openbabel as ob
# except ImportError:
#     print("Cannot find OpenBabel! Thus, bonds must be provided by the user.")
#     have_babel = False

"""
This module defines tools to calculate bond dissocation energies.
"""

__author__ = "Samuel Blau"
__copyright__ = "Copyright 2018, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Samuel Blau"
__email__ = "samblau1@gmail.com"
__status__ = "Alpha"
__date__ = "7/26/18"

logger = logging.getLogger(__name__)

class BondDissociationEnergies(MSONable):

    def __init__(self, molecule_entry, fragment_entries):
        """
        Standard constructor for bond dissociation energies.

        Args:

        Get list of all bonds in target molecule
        Build molecule graph of target molecule
        Loop through bonds aka edges and attempt to split into disconnected subgraphs
            On success:
                Search through fragments: find entries with initial OR final molecules which are isomorphic
                Should be three entires per fragment: original charge, OC+1, OC-1
                If there are too few entries:
                    If one of the fragments is an H atom, you should only find two. H+ should be set (to what value???)
                    If fragment matches initial, but not final structure:
                        Just accept the rearrangement as part of the energetic cost of dissociation?
                        At the very least, flag this structure change in the eventual entry
                    If fragment matches final, but not initial structure:
                        Ignore unless no other fragment matches either?
                    If fragment is unstable:
                        Find its subfragments? Is that even sensical?
                        Just pass?
                    If fragment truly cannot be found, aka has not been calculated or calculate failed to complete:
                        Tell user to calculate it and pass
                If there are too many entries:
                    Get rid of duplicates
                        Going to be annoying given time stamps, rounding errors, etc
                    Remove those whose initial rems don't conform to current defaults - Happens before passing
                    If still more than three, for now, isolate differences, report to user, and grab most recent
                    Just take more negative?
                Grab energies of each fragment at all three charges
                If target molecule is neutral:
                    Calculate E_frag1(0)  + E_frag2(0)
                    Calculate E_frag1(+1) + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(+1)
                If target molecule is +1:
                    Calculate E_frag1(0)  + E_frag2(+1)
                    Calculate E_frag1(+1) + E_frag2(0)
                    Calculate E_frag1(+2) + E_frag2(-1)?
                    Calculate E_frag1(-1) + E_frag2(+2)?
                If target molecule is -1:
                    Calculate E_frag1(0)  + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(0)
                    Calculate E_frag1(-2) + E_frag2(+1)?
                    Calculate E_frag1(+1) + E_frag2(-2)?
                Don't limit - just make final charge add to original charge
                Save E_molecule - calculated values associated with bond indices & charges in some data structure
            On failure:
                Track bonds that do not break to yield subgraphs as ring bonds
        Make all possible pairs of ring bonds
        Loop through ring bond pairs and attempt to split into disconnected subgraphs
            Same success/failure procedure as above. Should be wrapped in a function.
        Resolve individual rings
        Attempt to solve for single ring bond dissociation energies, careful with charge

        """

        self.molecule_entry = molecule_entry
        self.fragment_entries = fragment_entries
        self.bond_dissociation_energies = []
        # self.all_bonds = edges_from_babel(target_molecule)
        self.ring_bonds = []
        self.bad_pairs = []
        self.mol_graph = build_MoleculeGraph(Molecule.from_dict(molecule_entry["output"]["optimized_molecule"]))#, self.all_bonds)
        for bond in self.mol_graph.graph.edges: #switch to all_bonds?
            bonds = [(bond[0],bond[1])]
            self.fragment_and_process(bonds)
        self.bond_pairs = []
        for ii,bond in enumerate(self.ring_bonds):
            for jj in range(ii+1,len(self.ring_bonds)):
                bond_pair = [bond, self.ring_bonds[jj]]
                self.bond_pairs += [bond_pair]
        if len(self.bond_pairs) > 0:
            print("Ring bonds detected!")
        for bond_pair in self.bond_pairs:
            self.fragment_and_process(bond_pair)
        self.solve_ring_bonds()
        # return self.bond_dissociation_energies

    def fragment_and_process(self, bonds):
        try:
            frags = self.mol_graph.split_molecule_subgraphs(bonds,allow_reverse=True)
            frag_success = True
        except RuntimeError:
            if len(bonds) == 1:
                self.ring_bonds += bonds
            elif len(bonds) == 2:
                self.bad_pairs += bonds
            else:
                print('No reason to try and break more than two bonds at once! Exiting...')
                raise ValueError
            frag_success = False
        if frag_success:
            frag1_entries = self.search_fragment_entries(frags[0])
            frag2_entries = self.search_fragment_entries(frags[1])
            if bonds == [(0, 14)] or bonds == [(6, 10)]:
                # print()
                # print(bonds)
                # print()
                # print(frag1_entries[1])
                # print()
                # print(frag1_entries[2])
                # print()
                # print(frag2_entries[1])
                # print()
                # print(frag2_entries[2])
                # print()
            for frag1 in frag1_entries[0]:
                for frag2 in frag2_entries[0]:
                    if frag1["calcs_reversed"][-1]["input"]["rem"]["method"] == frag2["calcs_reversed"][-1]["input"]["rem"]["method"] and frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag2["calcs_reversed"][-1]["input"]["rem"]["basis"]: # Add PCM check?
                        if frag1["output"]["initial_molecule"]["charge"] + frag2["output"]["initial_molecule"]["charge"] == self.molecule_entry["output"]["optimized_molecule"]["charge"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["method"] == frag1["calcs_reversed"][-1]["input"]["rem"]["method"]:
                            coords = nx.get_node_attributes(self.mol_graph.graph, "coords")
                            specie = nx.get_node_attributes(self.mol_graph.graph, "specie")
                            if frag1["output"]["final_energy"] == "unstable" or frag2["output"]["final_energy"] == "unstable":
                                new_entry = ["unstable", bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], Molecule.from_dict(frag1["output"]["initial_molecule"]).composition.reduced_formula, frag1["output"]["initial_molecule"]["charge"], frag1["output"]["final_energy"], Molecule.from_dict(frag2["output"]["initial_molecule"]).composition.reduced_formula, frag2["output"]["initial_molecule"]["charge"], frag2["output"]["final_energy"]]
                            else:
                                new_entry = [self.molecule_entry["output"]["final_energy"] - (frag1["output"]["final_energy"] + frag2["output"]["final_energy"]), bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], Molecule.from_dict(frag1["output"]["initial_molecule"]).composition.reduced_formula, frag1["output"]["initial_molecule"]["charge"], frag1["output"]["final_energy"], Molecule.from_dict(frag2["output"]["initial_molecule"]).composition.reduced_formula, frag2["output"]["initial_molecule"]["charge"], frag2["output"]["final_energy"]]
                            self.bond_dissociation_energies += [new_entry]

    def search_fragment_entries(self, frag):
        entries = []
        initial_entries = []
        final_entries = []
        for entry in self.fragment_entries:
            initial_graph = build_MoleculeGraph(Molecule.from_dict(entry["input"]["initial_molecule"])).graph
            final_graph = build_MoleculeGraph(Molecule.from_dict(entry["output"]["initial_molecule"])).graph
            if is_isomorphic(frag.graph, initial_graph) and is_isomorphic(frag.graph, final_graph):
                entries += [entry]
            elif is_isomorphic(frag.graph, initial_graph):
                print(entry)
                initial_entries += [entry]
            elif is_isomorphic(frag.graph, final_graph):
                print(entry)
                final_entries += [entry]
        return [entries, initial_entries, final_entries]

    def solve_ring_bonds(self):
        pass
<commit_msg>debugging<commit_after># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

import re
import csv
import collections
import itertools
from io import open
import math
from six.moves import zip
import logging

from monty.json import MSONable, MontyDecoder

import numpy as np

from pymatgen.core.structure import Molecule
from pymatgen.io.qchem.outputs import edges_from_babel, build_MoleculeGraph, is_isomorphic
import networkx as nx

# have_babel = True
# try:
#     from pymatgen.io.babel import BabelMolAdaptor
#     import openbabel as ob
# except ImportError:
#     print("Cannot find OpenBabel! Thus, bonds must be provided by the user.")
#     have_babel = False

"""
This module defines tools to calculate bond dissocation energies.
"""

__author__ = "Samuel Blau"
__copyright__ = "Copyright 2018, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Samuel Blau"
__email__ = "samblau1@gmail.com"
__status__ = "Alpha"
__date__ = "7/26/18"

logger = logging.getLogger(__name__)

class BondDissociationEnergies(MSONable):

    def __init__(self, molecule_entry, fragment_entries):
        """
        Standard constructor for bond dissociation energies.

        Args:

        Get list of all bonds in target molecule
        Build molecule graph of target molecule
        Loop through bonds aka edges and attempt to split into disconnected subgraphs
            On success:
                Search through fragments: find entries with initial OR final molecules which are isomorphic
                Should be three entires per fragment: original charge, OC+1, OC-1
                If there are too few entries:
                    If one of the fragments is an H atom, you should only find two. H+ should be set (to what value???)
                    If fragment matches initial, but not final structure:
                        Just accept the rearrangement as part of the energetic cost of dissociation?
                        At the very least, flag this structure change in the eventual entry
                    If fragment matches final, but not initial structure:
                        Ignore unless no other fragment matches either?
                    If fragment is unstable:
                        Find its subfragments? Is that even sensical?
                        Just pass?
                    If fragment truly cannot be found, aka has not been calculated or calculate failed to complete:
                        Tell user to calculate it and pass
                If there are too many entries:
                    Get rid of duplicates
                        Going to be annoying given time stamps, rounding errors, etc
                    Remove those whose initial rems don't conform to current defaults - Happens before passing
                    If still more than three, for now, isolate differences, report to user, and grab most recent
                    Just take more negative?
                Grab energies of each fragment at all three charges
                If target molecule is neutral:
                    Calculate E_frag1(0)  + E_frag2(0)
                    Calculate E_frag1(+1) + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(+1)
                If target molecule is +1:
                    Calculate E_frag1(0)  + E_frag2(+1)
                    Calculate E_frag1(+1) + E_frag2(0)
                    Calculate E_frag1(+2) + E_frag2(-1)?
                    Calculate E_frag1(-1) + E_frag2(+2)?
                If target molecule is -1:
                    Calculate E_frag1(0)  + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(0)
                    Calculate E_frag1(-2) + E_frag2(+1)?
                    Calculate E_frag1(+1) + E_frag2(-2)?
                Don't limit - just make final charge add to original charge
                Save E_molecule - calculated values associated with bond indices & charges in some data structure
            On failure:
                Track bonds that do not break to yield subgraphs as ring bonds
        Make all possible pairs of ring bonds
        Loop through ring bond pairs and attempt to split into disconnected subgraphs
            Same success/failure procedure as above. Should be wrapped in a function.
        Resolve individual rings
        Attempt to solve for single ring bond dissociation energies, careful with charge

        """

        self.molecule_entry = molecule_entry
        self.fragment_entries = fragment_entries
        self.bond_dissociation_energies = []
        # self.all_bonds = edges_from_babel(target_molecule)
        self.ring_bonds = []
        self.bad_pairs = []
        self.mol_graph = build_MoleculeGraph(Molecule.from_dict(molecule_entry["output"]["optimized_molecule"]))#, self.all_bonds)
        for bond in self.mol_graph.graph.edges: #switch to all_bonds?
            bonds = [(bond[0],bond[1])]
            self.fragment_and_process(bonds)
        self.bond_pairs = []
        for ii,bond in enumerate(self.ring_bonds):
            for jj in range(ii+1,len(self.ring_bonds)):
                bond_pair = [bond, self.ring_bonds[jj]]
                self.bond_pairs += [bond_pair]
        if len(self.bond_pairs) > 0:
            print("Ring bonds detected!")
        for bond_pair in self.bond_pairs:
            self.fragment_and_process(bond_pair)
        self.solve_ring_bonds()
        # return self.bond_dissociation_energies

    def fragment_and_process(self, bonds):
        try:
            frags = self.mol_graph.split_molecule_subgraphs(bonds,allow_reverse=True)
            frag_success = True
        except RuntimeError:
            if len(bonds) == 1:
                self.ring_bonds += bonds
            elif len(bonds) == 2:
                self.bad_pairs += bonds
            else:
                print('No reason to try and break more than two bonds at once! Exiting...')
                raise ValueError
            frag_success = False
        if frag_success:
            frag1_entries = self.search_fragment_entries(frags[0])
            frag2_entries = self.search_fragment_entries(frags[1])
            # if bonds == [(0, 14)] or bonds == [(6, 10)]:
            #     print()
            #     print(bonds)
            #     print()
            #     print(frag1_entries[1])
            #     print()
            #     print(frag1_entries[2])
            #     print()
            #     print(frag2_entries[1])
            #     print()
            #     print(frag2_entries[2])
            #     print()
            for frag1 in frag1_entries[0]:
                for frag2 in frag2_entries[0]:
                    if frag1["calcs_reversed"][-1]["input"]["rem"]["method"] == frag2["calcs_reversed"][-1]["input"]["rem"]["method"] and frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag2["calcs_reversed"][-1]["input"]["rem"]["basis"]: # Add PCM check?
                        if frag1["output"]["initial_molecule"]["charge"] + frag2["output"]["initial_molecule"]["charge"] == self.molecule_entry["output"]["optimized_molecule"]["charge"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["method"] == frag1["calcs_reversed"][-1]["input"]["rem"]["method"]:
                            coords = nx.get_node_attributes(self.mol_graph.graph, "coords")
                            specie = nx.get_node_attributes(self.mol_graph.graph, "specie")
                            if frag1["output"]["final_energy"] == "unstable" or frag2["output"]["final_energy"] == "unstable":
                                new_entry = ["unstable", bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], Molecule.from_dict(frag1["output"]["initial_molecule"]).composition.reduced_formula, frag1["output"]["initial_molecule"]["charge"], frag1["output"]["final_energy"], Molecule.from_dict(frag2["output"]["initial_molecule"]).composition.reduced_formula, frag2["output"]["initial_molecule"]["charge"], frag2["output"]["final_energy"]]
                            else:
                                new_entry = [self.molecule_entry["output"]["final_energy"] - (frag1["output"]["final_energy"] + frag2["output"]["final_energy"]), bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], Molecule.from_dict(frag1["output"]["initial_molecule"]).composition.reduced_formula, frag1["output"]["initial_molecule"]["charge"], frag1["output"]["final_energy"], Molecule.from_dict(frag2["output"]["initial_molecule"]).composition.reduced_formula, frag2["output"]["initial_molecule"]["charge"], frag2["output"]["final_energy"]]
                            self.bond_dissociation_energies += [new_entry]

    def search_fragment_entries(self, frag):
        entries = []
        initial_entries = []
        final_entries = []
        for entry in self.fragment_entries:
            initial_graph = build_MoleculeGraph(Molecule.from_dict(entry["input"]["initial_molecule"])).graph
            final_graph = build_MoleculeGraph(Molecule.from_dict(entry["output"]["initial_molecule"])).graph
            if is_isomorphic(frag.graph, initial_graph) and is_isomorphic(frag.graph, final_graph):
                entries += [entry]
            elif is_isomorphic(frag.graph, initial_graph):
                print(entry)
                initial_entries += [entry]
            elif is_isomorphic(frag.graph, final_graph):
                print(entry)
                final_entries += [entry]
        return [entries, initial_entries, final_entries]

    def solve_ring_bonds(self):
        pass
<|endoftext|>