<commit_before>"""
Handles a single chunk of data (16x16x128 blocks) from a Minecraft save.
Chunk is currently McRegion only.
"""
from io import BytesIO
from struct import pack, unpack
import array, math

class Chunk(object):
	"""Class for representing a single chunk."""
	def __init__(self, nbt):
		chunk_data = nbt['Level']
		self.coords = chunk_data['xPos'],chunk_data['zPos']
		self.blocks = BlockArray(chunk_data['Blocks'].value, chunk_data['Data'].value)

	def get_coords(self):
		"""Return the coordinates of this chunk."""
		return (self.coords[0].value,self.coords[1].value)

	def __repr__(self):
		"""Return a representation of this Chunk."""
		return "Chunk("+str(self.coords[0])+","+str(self.coords[1])+")"

class BlockArray(object):
	"""Convenience class for dealing with a Block/data byte array."""
	def __init__(self, blocksBytes=None, dataBytes=None):
		"""Create a new BlockArray, defaulting to no block or data bytes."""
		if isinstance(blocksBytes, (bytearray, array.array)):
			self.blocksList = list(blocksBytes)
		else:
			self.blocksList = [0]*32768 # Create an empty block list (32768 entries of zero (air))

		if isinstance(dataBytes, (bytearray, array.array)):
			self.dataList = list(dataBytes)
		else:
			self.dataList = [0]*16384 # Create an empty data list (32768 4-bit entries of zero make 16384 byte entries)

	# Get all block entries
	def get_all_blocks(self):
		"""Return the blocks that are in this BlockArray."""
		return self.blocksList

	# Get all data entries
	def get_all_data(self):
		"""Return the data of all the blocks in this BlockArray."""
		bits = []
		for b in self.dataList:
			# The first byte of the Blocks arrays correspond
			# to the LEAST significant bits of the first byte of the Data.
			# NOT to the MOST significant bits, as you might expected.
			bits.append(b & 15) # Little end of the byte
			bits.append((b >> 4) & 15) # Big end of the byte
		return bits

	# Get all block entries and data entries as tuples
	def get_all_blocks_and_data(self):
		"""Return both blocks and data, packed together as tuples."""
		return list(zip(self.get_all_blocks(), self.get_all_data()))

	def get_blocks_struct(self):
		"""Return a dictionary with block ids keyed to (x, y, z)."""
		cur_x = 0
		cur_y = 0
		cur_z = 0
		blocks = {}
		for block_id in self.blocksList:
			blocks[(cur_x,cur_y,cur_z)] = block_id
			cur_y += 1
			if (cur_y > 127):
				cur_y = 0
				cur_z += 1
				if (cur_z > 15):
					cur_z = 0
					cur_x += 1
		return blocks

	# Give blockList back as a byte array
	def get_blocks_byte_array(self, buffer=False):
		"""Return a list of all blocks in this chunk."""
		if buffer:
			length = len(self.blocksList)
			return BytesIO(pack(">i", length)+self.get_blocks_byte_array())
		else:
			return array.array('B', self.blocksList).tostring()

	def get_data_byte_array(self, buffer=False):
		"""Return a list of data for all blocks in this chunk."""
		if buffer:
			length = len(self.dataList)
			return BytesIO(pack(">i", length)+self.get_data_byte_array())
		else:
			return array.array('B', self.dataList).tostring()

	def generate_heightmap(self, buffer=False, as_array=False):
		"""Return a heightmap, representing the highest solid blocks in this chunk."""
		non_solids = [0, 8, 9, 10, 11, 38, 37, 32, 31]
		if buffer:
			return BytesIO(pack(">i", 256)+self.generate_heightmap()) # Length + Heightmap, ready for insertion into Chunk NBT
		else:
			bytes = []
			for z in range(16):
				for x in range(16):
					for y in range(127, -1, -1):
						offset = y + z*128 + x*128*16
						if (self.blocksList[offset] not in non_solids or y == 0):
							bytes.append(y+1)
							break
			if (as_array):
				return bytes
			else:
				return array.array('B', bytes).tostring()

	def set_blocks(self, list=None, dict=None, fill_air=False):
		"""
		Sets all blocks in this chunk, using either a list or dictionary.
		Blocks not explicitly set can be filled to air by setting fill_air to True.
		"""
		if list:
			# Inputting a list like self.blocksList
			self.blocksList = list
		elif dict:
			# Inputting a dictionary like result of self.get_blocks_struct()
			list = []
			for x in range(16):
				for z in range(16):
					for y in range(128):
						coord = x,y,z
						offset = y + z*128 + x*128*16
						if (coord in dict):
							list.append(dict[coord])
						else:
							if (self.blocksList[offset] and not fill_air):
								list.append(self.blocksList[offset])
							else:
								list.append(0) # Air
			self.blocksList = list
		else:
			# None of the above...
			return False
		return True

	def set_block(self, x,y,z, id, data=0):
		"""Sets the block a x, y, z to the specified id, and optionally data."""
		offset = y + z*128 + x*128*16
		self.blocksList[offset] = id
		if (offset % 2 == 1):
			# offset is odd
			index = (offset-1)//2
			b = self.dataList[index]
			self.dataList[index] = (b & 240) + (data & 15) # modify lower bits, leaving higher bits in place
		else:
			# offset is even
			index = offset//2
			b = self.dataList[index]
			self.dataList[index] = (b & 15) + (data << 4 & 240) # modify ligher bits, leaving lower bits in place

	# Get a given X,Y,Z or a tuple of three coordinates
	def get_block(self, x,y,z, coord=False):
		"""Return the id of the block at x, y, z."""
		"""
		Laid out like:
		(0,0,0), (0,1,0), (0,2,0) ... (0,127,0), (0,0,1), (0,1,1), (0,2,1) ... (0,127,1), (0,0,2) ... (0,127,15), (1,0,0), (1,1,0) ... (15,127,15)

		::

		  blocks = []
		  for x in range(15):
		    for z in range(15):
		      for y in range(127):
		        blocks.append(Block(x,y,z))
		"""

		offset = y + z*128 + x*128*16 if (coord == False) else coord[1] + coord[2]*128 + coord[0]*128*16
		return self.blocksList[offset]

	# Get a given X,Y,Z or a tuple of three coordinates
	def get_data(self, x,y,z, coord=False):
		"""Return the data of the block at x, y, z."""
		offset = y + z*128 + x*128*16 if (coord == False) else coord[1] + coord[2]*128 + coord[0]*128*16
		# The first byte of the Blocks arrays correspond
		# to the LEAST significant bits of the first byte of the Data.
		# NOT to the MOST significant bits, as you might expected.
		if (offset % 2 == 1):
			# offset is odd
			index = (offset-1)//2
			b = self.dataList[index]
			return b & 15 # Get little (last 4 bits) end of byte
		else:
			# offset is even
			index = offset//2
			b = self.dataList[index]
			return (b >> 4) & 15 # Get big end (first 4 bits) of byte

	def get_block_and_data(self, x,y,z, coord=False):
		"""Return the tuple of (id, data) for the block at x, y, z"""
		return (self.get_block(x,y,z,coord),self.get_data(x,y,z,coord))

<commit_msg>Patch chunk.py to use Anvil format instead of MCRegion (not backward compatible currently).<commit_after>"""
Handles a single chunk of data (16x16x128 blocks) from a Minecraft save.
"""
from io import BytesIO
from struct import pack, unpack
import array, math

class Chunk(object):
	"""Class for representing a single chunk."""
	def __init__(self, nbt):
		chunk_data = nbt['Level']
		self.coords = chunk_data['xPos'],chunk_data['zPos']
		self.sections = [None]*16
		for section_data in chunk_data['Sections']:
		  self.sections[section_data['Y'].value] = BlockArray(section_data['Blocks'].value, section_data['Data'].value)

	def get_coords(self):
		"""Return the coordinates of this chunk."""
		return (self.coords[0].value,self.coords[1].value)

	def get_block(self, x, y, z):
		d,m = divmod(y, 16)
		if self.sections[d]:
			return self.sections[d].get_block(x, m, z)

	def __repr__(self):
		"""Return a representation of this Chunk."""
		return "Chunk("+str(self.coords[0])+","+str(self.coords[1])+")"

class BlockArray(object):
	"""Convenience class for dealing with a Block/data byte array."""
	def __init__(self, blocksBytes=None, dataBytes=None):
		"""Create a new BlockArray, defaulting to no block or data bytes."""
		if isinstance(blocksBytes, (bytearray, array.array)):
			self.blocksList = list(blocksBytes)
		else:
			self.blocksList = [0]*32768 # Create an empty block list (32768 entries of zero (air))

		if isinstance(dataBytes, (bytearray, array.array)):
			self.dataList = list(dataBytes)
		else:
			self.dataList = [0]*16384 # Create an empty data list (32768 4-bit entries of zero make 16384 byte entries)

	# Get all block entries
	def get_all_blocks(self):
		"""Return the blocks that are in this BlockArray."""
		return self.blocksList

	# Get all data entries
	def get_all_data(self):
		"""Return the data of all the blocks in this BlockArray."""
		bits = []
		for b in self.dataList:
			# The first byte of the Blocks arrays correspond
			# to the LEAST significant bits of the first byte of the Data.
			# NOT to the MOST significant bits, as you might expected.
			bits.append(b & 15) # Little end of the byte
			bits.append((b >> 4) & 15) # Big end of the byte
		return bits

	# Get all block entries and data entries as tuples
	def get_all_blocks_and_data(self):
		"""Return both blocks and data, packed together as tuples."""
		return list(zip(self.get_all_blocks(), self.get_all_data()))

	def get_blocks_struct(self):
		"""Return a dictionary with block ids keyed to (x, y, z)."""
		cur_x = 0
		cur_y = 0
		cur_z = 0
		blocks = {}
		for block_id in self.blocksList:
			blocks[(cur_x,cur_y,cur_z)] = block_id
			cur_y += 1
			if (cur_y > 127):
				cur_y = 0
				cur_z += 1
				if (cur_z > 15):
					cur_z = 0
					cur_x += 1
		return blocks

	# Give blockList back as a byte array
	def get_blocks_byte_array(self, buffer=False):
		"""Return a list of all blocks in this chunk."""
		if buffer:
			length = len(self.blocksList)
			return BytesIO(pack(">i", length)+self.get_blocks_byte_array())
		else:
			return array.array('B', self.blocksList).tostring()

	def get_data_byte_array(self, buffer=False):
		"""Return a list of data for all blocks in this chunk."""
		if buffer:
			length = len(self.dataList)
			return BytesIO(pack(">i", length)+self.get_data_byte_array())
		else:
			return array.array('B', self.dataList).tostring()

	def generate_heightmap(self, buffer=False, as_array=False):
		"""Return a heightmap, representing the highest solid blocks in this chunk."""
		non_solids = [0, 8, 9, 10, 11, 38, 37, 32, 31]
		if buffer:
			return BytesIO(pack(">i", 256)+self.generate_heightmap()) # Length + Heightmap, ready for insertion into Chunk NBT
		else:
			bytes = []
			for z in range(16):
				for x in range(16):
					for y in range(127, -1, -1):
						offset = y + z*128 + x*128*16
						if (self.blocksList[offset] not in non_solids or y == 0):
							bytes.append(y+1)
							break
			if (as_array):
				return bytes
			else:
				return array.array('B', bytes).tostring()

	def set_blocks(self, list=None, dict=None, fill_air=False):
		"""
		Sets all blocks in this chunk, using either a list or dictionary.
		Blocks not explicitly set can be filled to air by setting fill_air to True.
		"""
		if list:
			# Inputting a list like self.blocksList
			self.blocksList = list
		elif dict:
			# Inputting a dictionary like result of self.get_blocks_struct()
			list = []
			for x in range(16):
				for z in range(16):
					for y in range(128):
						coord = x,y,z
						offset = y + z*128 + x*128*16
						if (coord in dict):
							list.append(dict[coord])
						else:
							if (self.blocksList[offset] and not fill_air):
								list.append(self.blocksList[offset])
							else:
								list.append(0) # Air
			self.blocksList = list
		else:
			# None of the above...
			return False
		return True

	def set_block(self, x,y,z, id, data=0):
		"""Sets the block a x, y, z to the specified id, and optionally data."""
		offset = y + z*128 + x*128*16
		self.blocksList[offset] = id
		if (offset % 2 == 1):
			# offset is odd
			index = (offset-1)//2
			b = self.dataList[index]
			self.dataList[index] = (b & 240) + (data & 15) # modify lower bits, leaving higher bits in place
		else:
			# offset is even
			index = offset//2
			b = self.dataList[index]
			self.dataList[index] = (b & 15) + (data << 4 & 240) # modify ligher bits, leaving lower bits in place

	# Get a given X,Y,Z or a tuple of three coordinates
	def get_block(self, x,y,z, coord=False):
		"""Return the id of the block at x, y, z."""
		"""
		Laid out like:
		(0,0,0), (0,1,0), (0,2,0) ... (0,127,0), (0,0,1), (0,1,1), (0,2,1) ... (0,127,1), (0,0,2) ... (0,127,15), (1,0,0), (1,1,0) ... (15,127,15)

		::

		  blocks = []
		  for x in range(15):
		    for z in range(15):
		      for y in range(127):
		        blocks.append(Block(x,y,z))
		"""

		offset = y*16*16 + z*16 + x if (coord == False) else coord[0] + coord[2]*16 + coord[1]*16*16
		return self.blocksList[offset]

	# Get a given X,Y,Z or a tuple of three coordinates
	def get_data(self, x,y,z, coord=False):
		"""Return the data of the block at x, y, z."""
		offset = y*16*16 + z*16 + x if (coord == False) else coord[0] + coord[2]*16 + coord[1]*16*16
		# The first byte of the Blocks arrays correspond
		# to the LEAST significant bits of the first byte of the Data.
		# NOT to the MOST significant bits, as you might expected.
		if (offset % 2 == 1):
			# offset is odd
			index = (offset-1)//2
			b = self.dataList[index]
			return b & 15 # Get little (last 4 bits) end of byte
		else:
			# offset is even
			index = offset//2
			b = self.dataList[index]
			return (b >> 4) & 15 # Get big end (first 4 bits) of byte

	def get_block_and_data(self, x,y,z, coord=False):
		"""Return the tuple of (id, data) for the block at x, y, z"""
		return (self.get_block(x,y,z,coord),self.get_data(x,y,z,coord))

<|endoftext|>


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

from novaclient import exceptions as nova_exceptions

from trove.cluster import models
from trove.cluster.tasks import ClusterTasks
from trove.cluster.views import ClusterView
from trove.common import cfg
from trove.common import exception
from trove.common.i18n import _
from trove.common import remote
from trove.common.strategies.cluster import base
from trove.common import utils
from trove.common import wsgi
from trove.datastore import models as datastore_models
from trove.extensions.mgmt.clusters.views import MgmtClusterView
from trove.instance import models as inst_models
from trove.openstack.common import log as logging
from trove.quota.quota import check_quotas
from trove.taskmanager import api as task_api

LOG = logging.getLogger(__name__)
CONF = cfg.CONF

class MongoDbAPIStrategy(base.BaseAPIStrategy):

    @property
    def cluster_class(self):
        return MongoDbCluster

    @property
    def cluster_controller_actions(self):
        return {'add_shard': self._action_add_shard}

    def _action_add_shard(self, cluster, body):
        cluster.add_shard()
        return wsgi.Result(None, 202)

    @property
    def cluster_view_class(self):
        return MongoDbClusterView

    @property
    def mgmt_cluster_view_class(self):
        return MongoDbMgmtClusterView

class MongoDbCluster(models.Cluster):

    @classmethod
    def create(cls, context, name, datastore, datastore_version, instances):

        # TODO(amcreynolds): consider moving into CONF and even supporting
        # TODO(amcreynolds): an array of values, e.g. [3, 5, 7]
        # TODO(amcreynolds): or introduce a min/max num_instances and set
        # TODO(amcreynolds): both to 3
        num_instances = len(instances)
        if num_instances != 3:
            raise exception.ClusterNumInstancesNotSupported(num_instances=3)

        flavor_ids = [instance['flavor_id'] for instance in instances]
        if len(set(flavor_ids)) != 1:
            raise exception.ClusterFlavorsNotEqual()
        flavor_id = flavor_ids[0]
        nova_client = remote.create_nova_client(context)
        try:
            flavor = nova_client.flavors.get(flavor_id)
        except nova_exceptions.NotFound:
            raise exception.FlavorNotFound(uuid=flavor_id)
        mongo_conf = CONF.get(datastore_version.manager)
        num_configsvr = mongo_conf.num_config_servers_per_cluster
        num_mongos = mongo_conf.num_query_routers_per_cluster
        delta_instances = num_instances + num_configsvr + num_mongos
        deltas = {'instances': delta_instances}

        volume_sizes = [instance['volume_size'] for instance in instances
                        if instance.get('volume_size', None)]
        volume_size = None
        if mongo_conf.volume_support:
            if len(volume_sizes) != num_instances:
                raise exception.ClusterVolumeSizeRequired()
            if len(set(volume_sizes)) != 1:
                raise exception.ClusterVolumeSizesNotEqual()
            volume_size = volume_sizes[0]
            models.validate_volume_size(volume_size)
            # TODO(amcreynolds): for now, mongos+configsvr same flavor+disk
            deltas['volumes'] = volume_size * delta_instances
        else:
            # TODO(amcreynolds): is ephemeral possible for mongodb clusters?
            if len(volume_sizes) > 0:
                raise exception.VolumeNotSupported()
            ephemeral_support = mongo_conf.device_path
            if ephemeral_support and flavor.ephemeral == 0:
                raise exception.LocalStorageNotSpecified(flavor=flavor_id)

        check_quotas(context.tenant, deltas)

        db_info = models.DBCluster.create(
            name=name, tenant_id=context.tenant,
            datastore_version_id=datastore_version.id,
            task_status=ClusterTasks.BUILDING_INITIAL)

        replica_set_name = "rs1"

        member_config = {"id": db_info.id,
                         "shard_id": utils.generate_uuid(),
                         "instance_type": "member",
                         "replica_set_name": replica_set_name}
        for i in range(1, num_instances + 1):
            instance_name = "%s-%s-%s" % (name, replica_set_name, str(i))
            inst_models.Instance.create(context, instance_name,
                                        flavor_id,
                                        datastore_version.image_id,
                                        [], [], datastore,
                                        datastore_version,
                                        volume_size, None,
                                        availability_zone=None,
                                        nics=None,
                                        configuration_id=None,
                                        cluster_config=member_config)

        configsvr_config = {"id": db_info.id,
                            "instance_type": "config_server"}
        for i in range(1, num_configsvr + 1):
            instance_name = "%s-%s-%s" % (name, "configsvr", str(i))
            inst_models.Instance.create(context, instance_name,
                                        flavor_id,
                                        datastore_version.image_id,
                                        [], [], datastore,
                                        datastore_version,
                                        volume_size, None,
                                        availability_zone=None,
                                        nics=None,
                                        configuration_id=None,
                                        cluster_config=configsvr_config)

        mongos_config = {"id": db_info.id,
                         "instance_type": "query_router"}
        for i in range(1, num_mongos + 1):
            instance_name = "%s-%s-%s" % (name, "mongos", str(i))
            inst_models.Instance.create(context, instance_name,
                                        flavor_id,
                                        datastore_version.image_id,
                                        [], [], datastore,
                                        datastore_version,
                                        volume_size, None,
                                        availability_zone=None,
                                        nics=None,
                                        configuration_id=None,
                                        cluster_config=mongos_config)

        task_api.load(context, datastore_version.manager).create_cluster(
            db_info.id)

        return MongoDbCluster(context, db_info, datastore, datastore_version)

    def add_shard(self):

        if self.db_info.task_status != ClusterTasks.NONE:
            current_task = self.db_info.task_status.name
            msg = _("This action cannot be performed on the cluster while "
                    "the current cluster task is '%s'.") % current_task
            LOG.error(msg)
            raise exception.UnprocessableEntity(msg)

        db_insts = inst_models.DBInstance.find_all(cluster_id=self.id,
                                                   type='member').all()
        num_unique_shards = len(set([db_inst.shard_id for db_inst
                                     in db_insts]))
        arbitrary_shard_id = db_insts[0].shard_id
        members_in_shard = [db_inst for db_inst in db_insts
                            if db_inst.shard_id == arbitrary_shard_id]
        num_members_per_shard = len(members_in_shard)
        a_member = inst_models.load_any_instance(self.context,
                                                 members_in_shard[0].id)
        deltas = {'instances': num_members_per_shard}
        volume_size = a_member.volume_size
        if volume_size:
            deltas['volumes'] = volume_size * num_members_per_shard
        check_quotas(self.context.tenant, deltas)
        new_replica_set_name = "rs" + str(num_unique_shards + 1)
        new_shard_id = utils.generate_uuid()
        member_config = {"id": self.id,
                         "shard_id": new_shard_id,
                         "instance_type": "member",
                         "replica_set_name": new_replica_set_name}
        for i in range(1, num_members_per_shard + 1):
            instance_name = "%s-%s-%s" % (self.name, new_replica_set_name,
                                          str(i))
            inst_models.Instance.create(self.context, instance_name,
                                        a_member.flavor_id,
                                        a_member.datastore_version.image_id,
                                        [], [], a_member.datastore,
                                        a_member.datastore_version,
                                        volume_size, None,
                                        availability_zone=None,
                                        nics=None,
                                        configuration_id=None,
                                        cluster_config=member_config)

        self.update_db(task_status=ClusterTasks.ADDING_SHARD)
        manager = (datastore_models.DatastoreVersion.
                   load_by_uuid(db_insts[0].datastore_version_id).manager)
        task_api.load(self.context, manager).mongodb_add_shard_cluster(
            self.id,
            new_shard_id,
            new_replica_set_name)

class MongoDbClusterView(ClusterView):

    def build_instances(self):
        return self._build_instances(['query_router'], ['member'])

class MongoDbMgmtClusterView(MgmtClusterView):

    def build_instances(self):
        return self._build_instances(['query_router'],
                                     ['config_server',
                                      'member',
                                      'query_router'])
<commit_msg>Support nics and AZ for MongoDB clusters<commit_after># Copyright 2014 eBay Software Foundation
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

from novaclient import exceptions as nova_exceptions

from trove.cluster import models
from trove.cluster.tasks import ClusterTasks
from trove.cluster.views import ClusterView
from trove.common import cfg
from trove.common import exception
from trove.common.i18n import _
from trove.common import remote
from trove.common.strategies.cluster import base
from trove.common import utils
from trove.common import wsgi
from trove.datastore import models as datastore_models
from trove.extensions.mgmt.clusters.views import MgmtClusterView
from trove.instance import models as inst_models
from trove.openstack.common import log as logging
from trove.quota.quota import check_quotas
from trove.taskmanager import api as task_api

LOG = logging.getLogger(__name__)
CONF = cfg.CONF

class MongoDbAPIStrategy(base.BaseAPIStrategy):

    @property
    def cluster_class(self):
        return MongoDbCluster

    @property
    def cluster_controller_actions(self):
        return {'add_shard': self._action_add_shard}

    def _action_add_shard(self, cluster, body):
        cluster.add_shard()
        return wsgi.Result(None, 202)

    @property
    def cluster_view_class(self):
        return MongoDbClusterView

    @property
    def mgmt_cluster_view_class(self):
        return MongoDbMgmtClusterView

class MongoDbCluster(models.Cluster):

    @classmethod
    def create(cls, context, name, datastore, datastore_version, instances):

        # TODO(amcreynolds): consider moving into CONF and even supporting
        # TODO(amcreynolds): an array of values, e.g. [3, 5, 7]
        # TODO(amcreynolds): or introduce a min/max num_instances and set
        # TODO(amcreynolds): both to 3
        num_instances = len(instances)
        if num_instances != 3:
            raise exception.ClusterNumInstancesNotSupported(num_instances=3)

        flavor_ids = [instance['flavor_id'] for instance in instances]
        if len(set(flavor_ids)) != 1:
            raise exception.ClusterFlavorsNotEqual()
        flavor_id = flavor_ids[0]
        nova_client = remote.create_nova_client(context)
        try:
            flavor = nova_client.flavors.get(flavor_id)
        except nova_exceptions.NotFound:
            raise exception.FlavorNotFound(uuid=flavor_id)
        mongo_conf = CONF.get(datastore_version.manager)
        num_configsvr = mongo_conf.num_config_servers_per_cluster
        num_mongos = mongo_conf.num_query_routers_per_cluster
        delta_instances = num_instances + num_configsvr + num_mongos
        deltas = {'instances': delta_instances}

        volume_sizes = [instance['volume_size'] for instance in instances
                        if instance.get('volume_size', None)]
        volume_size = None
        if mongo_conf.volume_support:
            if len(volume_sizes) != num_instances:
                raise exception.ClusterVolumeSizeRequired()
            if len(set(volume_sizes)) != 1:
                raise exception.ClusterVolumeSizesNotEqual()
            volume_size = volume_sizes[0]
            models.validate_volume_size(volume_size)
            # TODO(amcreynolds): for now, mongos+configsvr same flavor+disk
            deltas['volumes'] = volume_size * delta_instances
        else:
            # TODO(amcreynolds): is ephemeral possible for mongodb clusters?
            if len(volume_sizes) > 0:
                raise exception.VolumeNotSupported()
            ephemeral_support = mongo_conf.device_path
            if ephemeral_support and flavor.ephemeral == 0:
                raise exception.LocalStorageNotSpecified(flavor=flavor_id)

        check_quotas(context.tenant, deltas)

        nics = [instance.get('nics', None) for instance in instances]

        azs = [instance.get('availability_zone', None)
               for instance in instances]

        db_info = models.DBCluster.create(
            name=name, tenant_id=context.tenant,
            datastore_version_id=datastore_version.id,
            task_status=ClusterTasks.BUILDING_INITIAL)

        replica_set_name = "rs1"

        member_config = {"id": db_info.id,
                         "shard_id": utils.generate_uuid(),
                         "instance_type": "member",
                         "replica_set_name": replica_set_name}
        for i in range(0, num_instances):
            instance_name = "%s-%s-%s" % (name, replica_set_name, str(i + 1))
            inst_models.Instance.create(context, instance_name,
                                        flavor_id,
                                        datastore_version.image_id,
                                        [], [], datastore,
                                        datastore_version,
                                        volume_size, None,
                                        availability_zone=azs[i],
                                        nics=nics[i],
                                        configuration_id=None,
                                        cluster_config=member_config)

        configsvr_config = {"id": db_info.id,
                            "instance_type": "config_server"}
        for i in range(1, num_configsvr + 1):
            instance_name = "%s-%s-%s" % (name, "configsvr", str(i))
            inst_models.Instance.create(context, instance_name,
                                        flavor_id,
                                        datastore_version.image_id,
                                        [], [], datastore,
                                        datastore_version,
                                        volume_size, None,
                                        availability_zone=None,
                                        nics=None,
                                        configuration_id=None,
                                        cluster_config=configsvr_config)

        mongos_config = {"id": db_info.id,
                         "instance_type": "query_router"}
        for i in range(1, num_mongos + 1):
            instance_name = "%s-%s-%s" % (name, "mongos", str(i))
            inst_models.Instance.create(context, instance_name,
                                        flavor_id,
                                        datastore_version.image_id,
                                        [], [], datastore,
                                        datastore_version,
                                        volume_size, None,
                                        availability_zone=None,
                                        nics=None,
                                        configuration_id=None,
                                        cluster_config=mongos_config)

        task_api.load(context, datastore_version.manager).create_cluster(
            db_info.id)

        return MongoDbCluster(context, db_info, datastore, datastore_version)

    def add_shard(self):

        if self.db_info.task_status != ClusterTasks.NONE:
            current_task = self.db_info.task_status.name
            msg = _("This action cannot be performed on the cluster while "
                    "the current cluster task is '%s'.") % current_task
            LOG.error(msg)
            raise exception.UnprocessableEntity(msg)

        db_insts = inst_models.DBInstance.find_all(cluster_id=self.id,
                                                   type='member').all()
        num_unique_shards = len(set([db_inst.shard_id for db_inst
                                     in db_insts]))
        arbitrary_shard_id = db_insts[0].shard_id
        members_in_shard = [db_inst for db_inst in db_insts
                            if db_inst.shard_id == arbitrary_shard_id]
        num_members_per_shard = len(members_in_shard)
        a_member = inst_models.load_any_instance(self.context,
                                                 members_in_shard[0].id)
        deltas = {'instances': num_members_per_shard}
        volume_size = a_member.volume_size
        if volume_size:
            deltas['volumes'] = volume_size * num_members_per_shard
        check_quotas(self.context.tenant, deltas)
        new_replica_set_name = "rs" + str(num_unique_shards + 1)
        new_shard_id = utils.generate_uuid()
        member_config = {"id": self.id,
                         "shard_id": new_shard_id,
                         "instance_type": "member",
                         "replica_set_name": new_replica_set_name}
        for i in range(1, num_members_per_shard + 1):
            instance_name = "%s-%s-%s" % (self.name, new_replica_set_name,
                                          str(i))
            inst_models.Instance.create(self.context, instance_name,
                                        a_member.flavor_id,
                                        a_member.datastore_version.image_id,
                                        [], [], a_member.datastore,
                                        a_member.datastore_version,
                                        volume_size, None,
                                        availability_zone=None,
                                        nics=None,
                                        configuration_id=None,
                                        cluster_config=member_config)

        self.update_db(task_status=ClusterTasks.ADDING_SHARD)
        manager = (datastore_models.DatastoreVersion.
                   load_by_uuid(db_insts[0].datastore_version_id).manager)
        task_api.load(self.context, manager).mongodb_add_shard_cluster(
            self.id,
            new_shard_id,
            new_replica_set_name)

class MongoDbClusterView(ClusterView):

    def build_instances(self):
        return self._build_instances(['query_router'], ['member'])

class MongoDbMgmtClusterView(MgmtClusterView):

    def build_instances(self):
        return self._build_instances(['query_router'],
                                     ['config_server',
                                      'member',
                                      'query_router'])
<|endoftext|>


<commit_before>#!/usr/bin/env python

try:
    import http.client as httplib
except ImportError:
    import httplib

import socket

from contextlib import contextmanager

@contextmanager
def tcpip4_socket(host, port):
    """Open a TCP/IP4 socket to designated host/port."""
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        s.connect((host, port))
        yield s
    finally:
        try:
            s.shutdown(socket.SHUT_RDWR)
        except OSError:
            pass
        finally:
            s.close()

@contextmanager
def http_connection(host, port):
    """Open an HTTP connection to designated host/port."""
    c = httplib.HTTPConnection(host, port)
    try:
       yield c
    finally:
       c.close()

<commit_msg>handle socket error<commit_after>#!/usr/bin/env python

try:
    import http.client as httplib
except ImportError:
    import httplib

import socket

from contextlib import contextmanager

@contextmanager
def tcpip4_socket(host, port):
    """Open a TCP/IP4 socket to designated host/port."""
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        s.connect((host, port))
        yield s
    finally:
        try:
            s.shutdown(socket.SHUT_RDWR)
+       except socket.error:
+           pass
        except OSError:
            pass
        finally:
            s.close()

@contextmanager
def http_connection(host, port):
    """Open an HTTP connection to designated host/port."""
    c = httplib.HTTPConnection(host, port)
    try:
       yield c
    finally:
       c.close()

<|endoftext|>


<commit_before>#!/usr/bin/python

import os
import csv
import json
import keyword
import sys
import re
import inflect
import djredcap
from django.core.management.base import BaseCommand, CommandError

header_keys = (
    'field_name',
    'form_name',
    'section_name',
    'field_type',
    'field_label',
    'choices',
    'field_note',
    'validation_type',
    'min_value',
    'max_value',
    'is_identifier',
    'branching_logic',
    'required',
    'custom_alignment',
    'question_number'
)
class Command(BaseCommand):
    help = """Attempts to read a REDCap data dictionary (CSV) and output a matching JSON file.
    Then attempts to read a JSON file and output matching Django models. Can take either a REDCap
    CSV file or a json file as input."""
    requires_model_validation = False;
    db_module = 'django.db'
    args = 'filename';

    def handle(self, fileName=None, *args, **options):
        if not fileName:
            raise CommandError('Enter a filename');

        fin = open(fileName);
        dialect = csv.Sniffer().sniff(fin.read(1024));
        fin.seek(0);
        reader = csv.DictReader(fin, fieldnames=header_keys,dialect=dialect);

        reader.next();

        fileName = djredcap.csv2json(self,reader, fileName);
        fileName1 = fileName;
        djredcap.json2dj(self,fileName);
        #os.remove(fileName1);
        os.remove(os.path.join(os.path.dirname(fileName1), fileName1));
<commit_msg>Added -JSON=jsonfile option to convert<commit_after>#!/usr/bin/python

import os
import csv
import json
import keyword
import sys
import re
import inflect
from optparse import make_option
import djredcap
from django.core.management.base import BaseCommand, CommandError

header_keys = (
    'field_name',
    'form_name',
    'section_name',
    'field_type',
    'field_label',
    'choices',
    'field_note',
    'validation_type',
    'min_value',
    'max_value',
    'is_identifier',
    'branching_logic',
    'required',
    'custom_alignment',
    'question_number'
)
class Command(BaseCommand):
    help = """Attempts to read a REDCap data dictionary (CSV) and output a matching JSON file.
    Then attempts to read a JSON file and output matching Django models. Can take either a REDCap
    CSV file or a json file as input."""
    requires_model_validation = False;
    db_module = 'django.db'
    args = 'filename';

    jsonFilename = '';

    option_list = BaseCommand.option_list + (
            make_option('--JSON', action='store',dest='json',
                help='Print a JSON file when using convert, default no'),
    )

    def handle(self, fileName=None, *args, **options):
        if not fileName:
            raise CommandError('Enter a filename');

        jsonFilename = options.get('json')
        fin = open(fileName);
        dialect = csv.Sniffer().sniff(fin.read(1024));
        fin.seek(0);
        reader = csv.DictReader(fin, fieldnames=header_keys,dialect=dialect);

        reader.next();

        fileName = djredcap.csv2json(self,reader, fileName);
        fileName1 = fileName;
        djredcap.json2dj(self,fileName);
        #os.remove(fileName1);
        if jsonFilename:
            os.rename(fileName1, jsonFilename)
        else:
            os.remove(os.path.join(os.path.dirname(fileName1), fileName1))
<|endoftext|>


<commit_before>from __future__ import print_function

import itertools

from snap import viewer

from snap.math import *
from snap.gl import *

import numpy as np

class Box(object):

    def __init__(self):
        self.size = np.ones(3)
        self.frame = Rigid3()

    def support(self, d):
        local = self.frame.orient.inv()(d)
        sign = (local > 0) - (local < 0)
        return self.frame(sign * self.size / 2)

def anderson(n, m, **kwargs):

    g = np.zeros( (n, m) )
    f = np.zeros( (n, m) )
    k = np.zeros( (m, m) )

    metric = kwargs.get('metric', np.ones(n) )

    old = np.zeros(n)
    delta = np.zeros(n)

    ones = np.ones(m)
    alpha = np.zeros(m)

    i = [0]

    def res(x):

        if x is None:
            g[:] = np.zeros( (n, m) )
            f[:] = np.zeros( (n, m) )
            k[:] = np.zeros( (m, m) )
            old[:] = np.zeros(n)
            delta[:] = np.zeros(n)
            return

        if i[0] > 0:
            delta[:] = x - old;

            index = i[0] % m

            g[:, index] = x
            f[:, index] = delta

            # TODO metric
            k[:, index] = f.transpose().dot( metric * f[:, index]  )
            k[index, :] = k[:, index].transpose()

            alpha = np.linalg.lstsq(k, ones)[0]

            s = sum(alpha)
            if s != 0:
                alpha /= s
                x[:] = g.dot(alpha)

        old[:] = x
        i[0] += 1

    return res

def pgs(M, q):

    n = q.size

    x = np.zeros(n)

    while True:

        for i in xrange(n):
            x[i] -= (M[i].dot(x) + q[i]) / M[i, i]
            x[i] = max(x[i], 0)

        yield x

def pgs_ls(A, b, p, omega = 1.0):
    '''M, q = A.dot(A.T), A.dot(p) - b'''

    m, n = A.shape

    x = np.zeros(m)

    # net = p + A.T.dot(x)
    net = np.copy(p)

    while True:

        for i in xrange(m):
            old = x[i]
            x[i] -= omega * (A[i].dot(net) - b[i]) / A[i].dot(A[i])
            x[i] = max(x[i], 0.0)
            # net += A[i] * ( x[i] - old )

            net[:] = p + A.T.dot(x)

        yield x, net

def pocs(A, b, p):

    m, n = A.shape

    I = np.zeros( (m, n) )

    q = np.copy(p)

    def proj(i, x):
        y = b[i] - A[i].dot(x)
        if y <= 0: return x
        res = x + A[i] * (y / (A[i].dot(A[i])))
        return res

    n = 0

    while True:
        i = n % m
        p =  q - I[i]

        q[:] = proj(i, p)
        I[i] = q - p

        yield q
        n += 1

def project_simplex(S, p):

    m, n = S.shape
    assert n + 1 == m

    Q = np.zeros( (m, m) )

    Q[:n] = S.T
    Q[-1] = 1

    Qinv = np.linalg.inv(Q)

    A = Qinv[:, :n]
    b = - Qinv[:, -1]

    res = []

    for it, (x, net) in itertools.izip(xrange(n * n), pgs_ls(A, b, p) ):
        res.append(np.copy(net))
    return res

    for it, x in itertools.izip(xrange(2 * n * n), pocs(A, b, p)):
        res.append(np.copy(x))

    return res

S = np.zeros( (4, 3) )
S[1:] = np.identity(3)

p = np.ones(3)
q = project_simplex(S, p)

def reset(): pass

def init(): pass

def keypress(key):
    if key == ' ':

        S[1:] = 2 * np.random.rand( 3, 3 ) - 1
        p[:] = 2 * np.random.rand(3) - 1
        q[:] = project_simplex(S, p)

def animate(): pass

def draw():
    glDisable(GL_LIGHTING)
    glPointSize(5)
    glLineWidth(2)
    glColor(1, 1, 1)

    glBegin(GL_POINTS)
    glVertex(p)

    for qi in q:
       glVertex(qi)

    glEnd()

    glBegin(GL_LINES)

    for i, j in itertools.product(xrange(4), xrange(4)):
        glVertex(S[i])
        glVertex(S[j])

    glEnd()

    glBegin(GL_LINE_STRIP)
    glVertex(p)

    for qi in q:
        glVertex(qi)

    glEnd()

    glEnable(GL_LIGHTING)

viewer.run()
<commit_msg>proj splx<commit_after>from __future__ import print_function

import itertools

from snap import viewer

from snap.math import *
from snap.gl import *

import numpy as np

class Box(object):

    def __init__(self):
        self.size = np.ones(3)
        self.frame = Rigid3()

    def support(self, d):
        local = self.frame.orient.inv()(d)
        sign = (local > 0) - (local < 0)
        return self.frame(sign * self.size / 2)

def barycentric_coordinates(S, p):
    m, n = S.shape

    A = np.zeros( (n + 1, m))
    A[:n] = S.T
    A[-1] = 1
    b = np.zeros(m)

    b[:n] = p
    b[-1] = 1

    x = np.linalg.solve(A, b)
    return x

def pgs(M, q):

    n = q.size

    x = np.zeros(n)

    while True:

        for i in xrange(n):
            x[i] -= (M[i].dot(x) + q[i]) / M[i, i]
            x[i] = max(x[i], 0)

        yield x

def pgs_ls(A, b, p, omega = 1.0):
    '''M, q = A.dot(A.T), A.dot(p) - b'''

    m, n = A.shape

    x = np.zeros(m)

    # net = p + A.T.dot(x)
    net = np.copy(p)

    while True:

        for i in xrange(m):
            old = x[i]
            x[i] -= omega * (A[i].dot(net) - b[i]) / A[i].dot(A[i])
            x[i] = max(x[i], 0.0)
            # net += A[i] * ( x[i] - old )

            net[:] = p + A.T.dot(x)

        yield x, net

def pocs(A, b, p):

    m, n = A.shape

    I = np.zeros( (m, n) )

    q = np.copy(p)

    def proj(i, x):
        y = b[i] - A[i].dot(x)
        if y <= 0: return x
        res = x + A[i] * (y / (A[i].dot(A[i])))
        return res

    n = 0

    while True:
        i = n % m
        p =  q - I[i]

        q[:] = proj(i, p)
        I[i] = q - p

        yield q
        n += 1

def proj_splx(x):
    n = x.size
    y = sorted(x)

    s = 0
    for k in range(n):
        i = n - 1 - k
        s += y[i]
        t = (s - 1.0) / float(n - i)
        if t > y[i - 1]:
            break

    return np.maximum(x - t, 0)

def project_simplex(S, p):

    m, n = S.shape
    assert n + 1 == m

    Q = np.zeros( (m, m) )

    Q[:n] = S.T
    Q[-1] = 1

    Qinv = np.linalg.inv(Q)

    A = Qinv[:, :n]
    b = - Qinv[:, -1]

    res = []

    for it, (x, net) in itertools.izip(xrange(n * n), pgs_ls(A, b, p) ):
        res.append(np.copy(net))
    return res

    for it, x in itertools.izip(xrange(2 * n * n), pocs(A, b, p)):
        res.append(np.copy(x))

    return res

S = np.zeros( (4, 3) )
S[1:] = np.identity(3)

p = np.ones(3)
q = proj_splx(p) #project_simplex(S, p)

def reset(): pass

def init(): pass

def keypress(key):
    if key == ' ':

        # S[1:] = 2 * np.random.rand( 3, 3 ) - 1
        p[:] = 2 * np.random.rand(3) - 1

        q[:] = proj_splx(p)
        # q[:] = project_simplex(S, p)

def animate(): pass

def draw():
    glDisable(GL_LIGHTING)
    glPointSize(5)
    glLineWidth(2)
    glColor(1, 1, 1)

    glBegin(GL_POINTS)
    glVertex(p)
    glVertex(q)

    for v in xrange(3):
        col = np.zeros(3)
        col[v] = 1
        glColor(col)
        glVertex(S[1 + v])

    glEnd()

    glColor(1, 1, 1)
    glBegin(GL_LINES)

    for i, j in itertools.product(xrange(4), xrange(4)):
        glVertex(S[i])
        glVertex(S[j])

    glEnd()

    glBegin(GL_LINE_STRIP)
    glVertex(p)

    glVertex(q)

    glEnd()

    glEnable(GL_LIGHTING)

viewer.run()
<|endoftext|>


<commit_before>from __future__ import print_function, division
import  sys
sys.dont_write_bytecode = True
from table import *
from counts import *

@setting
def CUT(): return o(
  crowded=4,
  cohen=0.1,
  fayyad=False
)

def crowded(n):
  return n > the.CUT.crowded

def smallEffectSize(lst,num):
  return Num(num(z) for z in lst).sd()*the.CUT.cohen

def recurse(this, divisor, id, x,cuts):
    cut,about = divisor(this)
    if cut:
      recurse(this[:cut], divisor, id,x, cuts);
      recurse(this[cut:], divisor, id,x, cuts)
    else:
      cuts += [o( id = id,
                  n  = len(cuts),
                  x  = o(lo=x(this[0]), hi=x(this[-1])),
                  y  = about,
                  has= this)]
    return cuts

def spliters(this,lhs,rhs,x,y,small):
  def silly():
    #if j > the.CUT.crowded:
     # return True
    return x(this[j]) - x(this[0]) < small
    #if j + 1  < len(this):
     # if x(this[j+1]) -  x(this[j]) < small :
      #  return True
  for j,one in enumerate(this):
    if crowded(lhs.n):
      if crowded(rhs.n):
        if not silly():
          yield j,one
    rhs -= y(one)
    lhs += y(one)

def sdiv1(lst,x=None,**d):
  return sdiv(lst,num1=x,num2=x,**d)

def sdiv(lst, id=0, small=None,
         num1= lambda z:z[0],
         num2= lambda z:z[-1]):
  def sdivide(this):
    lhs,rhs = Num(), Num(num2(z) for z in this)
    n0, sd0, cut, mu = rhs.n, rhs.sd(), None, rhs.mu
    score = sd0
    for j,one  in spliters(this,lhs,rhs,num1,num2,small):
      maybe= lhs.n/n0*lhs.sd()+ rhs.n/n0*rhs.sd()
      if maybe < score:
        if abs(lhs.mu - rhs.mu) >= small:
          cut,score = j,maybe
    return cut, o(mu=mu,n=n0,sd=sd0)
  if not lst: return []
  small = small or smallEffectSize(lst,num1)
  return recurse(sorted(lst,key=num1),
                 sdivide, id, num1, [] )

def ediv(lst, id=0, small=None,
         num= lambda x:x[0],
         sym= lambda x:x[1]):
  def edivide(this):
    def ke(z): return z.k()*z.ent()
    lhs,rhs   = Sym(),Sym(sym(x) for x in this)
    n0,k0,e0,ke0= rhs.n,rhs.k(),rhs.ent(),ke(rhs)
    cut, least  = None, e0
    for j,one  in spliters(this,lhs,rhs,num,sym,small):
      maybe= lhs.n/n0*lhs.ent()+ rhs.n/n0*rhs.ent()
      if maybe < least :
        if the.CUT.fayyad:
          gain  = e0 - maybe
          delta = log2(3**k0 -2)- (ke0 - ke(rhs) - ke(lhs))
          if gain >= (log2(n0 - 1) + delta)/n0:
              cut,least = j,maybe
        else:
          cut,least = j,maybe
    return cut,o(n=n0,e=e0)
  if not lst: return []
  small = small or smallEffectSize(lst,num)
  return recurse(sorted(lst,key=num),
                 edivide, id, num, [])

t = table(cols(FILE('data/albrecht.csv')))

klasses= sdiv1(t.rows,  x= lambda z:z.raw[-1])

for row in t.rows:
  row.cooked = row.raw[:]
for klass in klasses:
  k = klass.n
  for row in klass.has:
    row.cooked[-1] =  k

for n in t.inNums:
  for r in ediv(t.rows,
                  num =lambda z:z.raw[n],
                  sym =lambda z:z.cooked[-1]):
    print(n,r.x,len(r.has),r.y)

#print(len(t.rows))
## print("\n",x)

def _ediv():
  "Demo code to test the above."
  import random
  bell= random.gauss
  random.seed(1)
  def go(lst):
    print(""); print(sorted(lst)[:],"...")
    d = ediv(lst)
    print(d[0][1][0][0])

    print(d[0] ,"d[0]")
    print(d[1] ,"d[1]")
    print(d[0][1], "d[0][1]")
    print(d[0][1][0], "d[0][1][0]")
    print(d[0][1][0][0], "d[0][1][0][0]")

    for d in  ediv(lst):
      print(d[1][0])
  X,Y="X","Y"
  l=[(1,X),(2,X),(3,X),(4,X),(11,Y),(12,Y),(13,Y),(14,Y)]
  #l=[(3.02,'x'),(3.05,'x'),(3.03,'x'),(4.00,'y'),(4.01,'x'),(4.02,'y')]
  go(l)
  """  go(l)
  l[0] = (1,Y)
  go(l)
  go(l*2)
  go([(1,X),(2,X),(3,X),(4,X),(11,X),(12,X),(13,X),(14,X)])

  go([(64,X),(65,Y),(68,X),(69,Y),(70,X),(71,Y),
      (72,X),(72,Y),(75,X),(75,X),
      (80,Y),(81,Y),(83,Y),(85,Y)]*2)
  l=[]
  for _ in range(1000):
    l += [(bell(20,1),  X),(bell(10,1),Y),
          (bell(30,1),'Z'),(bell(40,1),'W')]
  go(l)
  go([(1,X)])
  """
#__name__ == '__main__' and _ediv()

"""
Output:

[(1, 'X'), (2, 'X'), (3, 'X'), (4, 'X'),
 (11, 'Y'), (12, 'Y'), (13, 'Y'), (14, 'Y')] ...
1
11

[(1, 'Y'), (2, 'X'), (3, 'X'), (4, 'X'),
 (11, 'Y'), (12, 'Y'), (13, 'Y'), (14, 'Y')] ...
1

[(1, 'Y'), (1, 'Y'), (2, 'X'), (2, 'X'),
 (3, 'X'), (3, 'X'), (4, 'X'), (4, 'X'), (11, 'Y'), (11, 'Y')] ...
1
11

[(1, 'X'), (2, 'X'), (3, 'X'), (4, 'X'),
 (11, 'X'), (12, 'X'), (13, 'X'), (14, 'X')] ...
1

[(64, 'X'), (64, 'X'), (65, 'Y'), (65, 'Y'),
 (68, 'X'), (68, 'X'), (69, 'Y'), (69, 'Y'), (70, 'X'), (70, 'X')] ...
64
80

[(6.900378121061215, 'Y'), (7.038785729480842, 'Y'),
 (7.31690058848835, 'Y'), (7.359039915634471, 'Y'),
 (7.364480069138072, 'Y'), (7.553496312538384, 'Y'),
 (7.581606303196569, 'Y'), (7.651878578401048, 'Y'),
 (7.655341871448137, 'Y'), (7.677081766167625, 'Y')] ...
6.90037812106
16.907507693
26.850034984
36.8600218357

[(1, 'X')] ...
1

"""
<commit_msg>saving<commit_after>from __future__ import print_function, division
import  sys
sys.dont_write_bytecode = True
from table import *
from counts import *

@setting
def CUT(): return o(
  crowded=4,
  cohen=0.1,
  fayyad=False
)

def crowded(n):
  return n >= the.CUT.crowded

def smallEffectSize(lst,num):
  return Num(num(z) for z in lst).sd()*the.CUT.cohen

def recurse(this, divisor, id, x,cuts):
    cut,about = divisor(this)
    if cut:
      recurse(this[:cut], divisor, id,x, cuts);
      recurse(this[cut:], divisor, id,x, cuts)
    else:
      cuts += [o( id = id,
                  n  = len(cuts),
                  x  = o(lo=x(this[0]), hi=x(this[-1])),
                  y  = about,
                  has= this)]
    return cuts

def spliters(this,lhs,rhs,x,y,small):
  def silly():
    #if j > the.CUT.crowded:
     # return True
    return x(this[j]) - x(this[0]) < small
    #if j + 1  < len(this):
     # if x(this[j+1]) -  x(this[j]) < small :
      #  return True
  for j,one in enumerate(this):
    rhs -= y(one)
    lhs += y(one)
    if crowded(lhs.n):
      if crowded(rhs.n):
        if not silly():
          yield j+1,one

def sdiv1(lst,x=None,**d):
  return sdiv(lst,num1=x,num2=x,**d)

def sdiv(lst, id=0, small=None,
         num1= lambda z:z[0],
         num2= lambda z:z[-1]):
  def sdivide(this):
    lhs,rhs = Num(), Num(num2(z) for z in this)
    n0, sd0, cut, mu = rhs.n, rhs.sd(), None, rhs.mu
    score = sd0
    for j,one  in spliters(this,lhs,rhs,num1,num2,small):
      maybe= lhs.n/n0*lhs.sd()+ rhs.n/n0*rhs.sd()
      if maybe < score:
        if abs(lhs.mu - rhs.mu) >= small:
          cut,score = j,maybe
    return cut, o(mu=mu,n=n0,sd=sd0)
  if not lst: return []
  small = small or smallEffectSize(lst,num1)
  return recurse(sorted(lst,key=num1),
                 sdivide, id, num1, [] )

def ediv(lst, id=0, small=None,
         num= lambda x:x[0],
         sym= lambda x:x[1]):
  def edivide(this):
    def ke(z): return z.k()*z.ent()
    lhs,rhs   = Sym(),Sym(sym(x) for x in this)
    n0,k0,e0,ke0= rhs.n,rhs.k(),rhs.ent(),ke(rhs)
    cut, least  = None, e0
    for j,one  in spliters(this,lhs,rhs,num,sym,small):
      maybe= lhs.n/n0*lhs.ent()+ rhs.n/n0*rhs.ent()
      if maybe < least :
        if the.CUT.fayyad:
          gain  = e0 - maybe
          delta = log2(3**k0 -2)- (ke0 - ke(rhs) - ke(lhs))
          if gain >= (log2(n0 - 1) + delta)/n0:
              cut,least = j,maybe
        else:
          cut,least = j,maybe
    return cut,o(n=n0,e=e0)
  if not lst: return []
  small = small or smallEffectSize(lst,num)
  return recurse(sorted(lst,key=num),
                 edivide, id, num, [])

#t = table(cols(FILE('data/albrecht.csv')))
t = table(cols(FILE('data/nasa93.csv')))

klasses= sdiv1(t.rows,  x= lambda z:z.raw[-1])

for row in t.rows:
  row.cooked = row.raw[:]
for klass in klasses:
  k = klass.n
  for row in klass.has:
    row.cooked[-1] =  k

for n in t.inNums:
  for r in ediv(t.rows,
                  num =lambda z:z.raw[n],
                  sym =lambda z:z.cooked[-1]):
    print(n,r.x,len(r.has),r.y)

#print(len(t.rows))
## print("\n",x)

def _ediv():
  "Demo code to test the above."
  import random
  bell= random.gauss
  random.seed(1)
  def go(lst):
    print(""); print(sorted(lst)[:],"...")
    d = ediv(lst)
    print(d[0][1][0][0])

    print(d[0] ,"d[0]")
    print(d[1] ,"d[1]")
    print(d[0][1], "d[0][1]")
    print(d[0][1][0], "d[0][1][0]")
    print(d[0][1][0][0], "d[0][1][0][0]")

    for d in  ediv(lst):
      print(d[1][0])
  X,Y="X","Y"
  l=[(1,X),(2,X),(3,X),(4,X),(11,Y),(12,Y),(13,Y),(14,Y)]
  #l=[(3.02,'x'),(3.05,'x'),(3.03,'x'),(4.00,'y'),(4.01,'x'),(4.02,'y')]
  go(l)
  """  go(l)
  l[0] = (1,Y)
  go(l)
  go(l*2)
  go([(1,X),(2,X),(3,X),(4,X),(11,X),(12,X),(13,X),(14,X)])

  go([(64,X),(65,Y),(68,X),(69,Y),(70,X),(71,Y),
      (72,X),(72,Y),(75,X),(75,X),
      (80,Y),(81,Y),(83,Y),(85,Y)]*2)
  l=[]
  for _ in range(1000):
    l += [(bell(20,1),  X),(bell(10,1),Y),
          (bell(30,1),'Z'),(bell(40,1),'W')]
  go(l)
  go([(1,X)])
  """
#__name__ == '__main__' and _ediv()

"""
Output:

[(1, 'X'), (2, 'X'), (3, 'X'), (4, 'X'),
 (11, 'Y'), (12, 'Y'), (13, 'Y'), (14, 'Y')] ...
1
11

[(1, 'Y'), (2, 'X'), (3, 'X'), (4, 'X'),
 (11, 'Y'), (12, 'Y'), (13, 'Y'), (14, 'Y')] ...
1

[(1, 'Y'), (1, 'Y'), (2, 'X'), (2, 'X'),
 (3, 'X'), (3, 'X'), (4, 'X'), (4, 'X'), (11, 'Y'), (11, 'Y')] ...
1
11

[(1, 'X'), (2, 'X'), (3, 'X'), (4, 'X'),
 (11, 'X'), (12, 'X'), (13, 'X'), (14, 'X')] ...
1

[(64, 'X'), (64, 'X'), (65, 'Y'), (65, 'Y'),
 (68, 'X'), (68, 'X'), (69, 'Y'), (69, 'Y'), (70, 'X'), (70, 'X')] ...
64
80

[(6.900378121061215, 'Y'), (7.038785729480842, 'Y'),
 (7.31690058848835, 'Y'), (7.359039915634471, 'Y'),
 (7.364480069138072, 'Y'), (7.553496312538384, 'Y'),
 (7.581606303196569, 'Y'), (7.651878578401048, 'Y'),
 (7.655341871448137, 'Y'), (7.677081766167625, 'Y')] ...
6.90037812106
16.907507693
26.850034984
36.8600218357

[(1, 'X')] ...
1

"""
<|endoftext|>


<commit_before>import logging
import pathlib
import tempfile
from collections import ChainMap

from tqdm import tqdm

from ppqm import chembridge, constants, linesio, shell, units
from ppqm.calculator import BaseCalculator

G16_CMD = "g16"
G16_FILENAME = "_tmp_g16_input.com"

_logger = logging.getLogger("g16")

COLUMN_SCF_ENERGY = "scf_energy"
COLUMN_CONVERGED = "is_converged"
COLUMN_MULIKEN_CHARGES = "mulliken charges"
COLUMN_CM5_CHARGES = "cm5_charges"
COLUMN_HIRSHFELD_CHARGES = "hirshfeld_charges"
COLUMN_NBO_BONDORDER = "bond_orders"
COLUMN_SHIELDING_CONSTANTS = "shielding_constants"

class GaussianCalculator(BaseCalculator):
    def __init__(
        self,
        cmd=G16_CMD,
        filename=G16_FILENAME,
        show_progress=False,
        n_cores=None,
        memory=2,
        **kwargs,
    ):

        super().__init__(**kwargs)

        self.cmd = cmd
        self.filename = filename
        self.n_cores = n_cores
        self.memory = memory
        self.show_progress = show_progress

        self.g16_options = {
            "cmd": self.cmd,
            "scr": self.scr,
            "filename": self.filename,
            "memory": self.memory,
        }

        # Default G16 options
        self.options = {}

        #
        self.health_check()

    def __repr__(self) -> str:
        return f"G16Calc(cmd={self.cmd}, scr={self.scr}, n_cores={self.n_cores}, memory={self.memory}gb)"

    def health_check(self):
        """ """

    def calculate(self, molobj, options, footer=None):
        """ """

        if self.n_cores > 1:
            raise NotImplementedError("Parallel not implemented yet.")
        else:
            results = self.calculate_serial(molobj, options, footer=footer)

        return results

    def calculate_serial(self, molobj, options, footer=None):
        """ """

        # If not singlet "spin" is part of options
        if "spin" in options.keys():
            spin = str(options.pop("spin"))
        else:
            spin = str(1)

        options_prime = ChainMap(options, self.options)
        options_prime = dict(options_prime)

        n_confs = molobj.GetNumConformers()
        atoms, _, charge = chembridge.get_axyzc(molobj, atomfmt=str)

        if self.show_progress:
            pbar = tqdm(
                total=n_confs,
                desc="g16(1)",
                **constants.TQDM_OPTIONS,
            )

        properties_list = []
        for conf_idx in range(n_confs):

            coord = chembridge.get_coordinates(molobj, confid=conf_idx)

            properties = get_properties_from_axyzc(
                atoms,
                coord,
                charge,
                spin,
                options=options_prime,
                footer=footer,
                **self.g16_options,
            )

            properties_list.append(properties)

            if self.show_progress:
                pbar.update(1)

        if self.show_progress:
            pbar.close()

        return properties_list

def get_properties_from_axyzc(
    atoms_str,
    coordinates,
    charge,
    spin,
    options=None,
    scr=constants.SCR,
    clean_tempfiles=False,
    cmd=G16_CMD,
    filename=G16_FILENAME,
    footer=None,
    **kwargs,
):

    if isinstance(scr, str):
        scr = pathlib.Path(scr)

    if not filename.endswith(".com"):
        filename += ".com"

    tempdir = tempfile.TemporaryDirectory(dir=scr, prefix="g16_")
    scr = pathlib.Path(tempdir.name)

    # write input file
    input_header = get_header(options, **kwargs)

    inputstr = get_inputfile(
        atoms_str, coordinates, charge, spin, input_header, footer=footer, title="g16 input"
    )
    print(inputstr)
    with open(scr / filename, "w") as f:
        f.write(inputstr)

    # Run subprocess cmd
    cmd = " ".join([cmd, filename])
    _logger.debug(cmd)

    lines = shell.stream(cmd, cwd=scr)
    lines = list(lines)

    # for line in lines:
    #    print(line.strip())
    # print("="*80 )
    # print()

    termination_pattern = "Normal termination of Gaussian"
    idx = linesio.get_rev_index(lines, termination_pattern, stoppattern="File lengths")
    if idx is None:
        _logger.critical("Abnormal termination of Gaussian")
        return None

    # Parse properties from Gaussian output
    properties = read_properties(lines, options)

    # Clean scr dir. TODO: Does this work??
    if clean_tempfiles:
        tempdir.cleanup()

    return properties

def get_inputfile(atom_strs, coordinates, charge, spin, header, footer=None, **kwargs):
    """ """

    inputstr = header + 2 * "\n"
    inputstr += f"  {kwargs['title']}" + 2 * "\n"

    inputstr += f"{charge}  {spin} \n"
    for atom_str, coord in zip(atom_strs, coordinates):
        inputstr += f"{atom_str}  " + " ".join([str(x) for x in coord]) + "\n"
    inputstr += "\n"  # magic line

    if footer is not None:
        inputstr += footer + "\n"

    return inputstr

def get_header(options, **kwargs):
    """ Write G16 header """

    header = f"%mem={kwargs.pop('memory')}gb\n"
    header += "# "
    for key, value in options.items():
        if (value is None) or (len(value) == 0):
            header += f"{key} "
        elif isinstance(value, str):
            header += f"{key}={value} "
        else:
            header += f"{key}=("
            for subkey, subvalue in value.items():
                if (subvalue is None) or (len(subvalue) == 0):
                    header += f"{subkey}, "
                else:
                    header += f"{subkey}={subvalue}, "
            header = header[:-2] + ") "

    return header

def read_properties(lines, options):
    """ Extract values from output depending on calculation options """

    # Collect readers
    readers = []

    if "opt" in options:
        raise NotImplementedError("not implemented opt properties parser")
        # reader = read_properties_opt
    else:
        readers.append(read_properties_sp)

    if "pop" in options:
        if "nboread" in options["pop"]:
            readers.append(get_nbo_bond_orders)

        if "hirsfeld" in options["pop"]:
            readers.append(get_hirsfeld_charges)
            readers.append(get_cm5_charges)

    if "nmr" in options:
        readers.append(get_nmr_shielding_constants)

    # Get properties
    properties = dict()
    for reader in readers:
        new_properties = reader(lines)
        assert isinstance(new_properties, dict)

        properties.update(new_properties)

    return properties

def read_properties_sp(lines):
    """
    Read Mulliken charges
    """
    properties = dict()

    for line in lines:
        if "SCF Done:  E(" in line:
            scf_energy = float(line.split()[4]) * units.hartree_to_kcalmol
            properties[COLUMN_SCF_ENERGY] = scf_energy
            break

    properties.update(get_mulliken_charges(lines))

    return properties

def read_properties_opt(lines):
    """ """

def get_mulliken_charges(lines):
    """ Read Mulliken charges """
    keywords = ["Mulliken charges:", "Sum of Mulliken charges"]
    start, stop = linesio.get_rev_indices_patterns(lines, keywords)
    mulliken_charges = [float(x.split()[-1]) for x in lines[start + 2 : stop]]
    return {COLUMN_MULIKEN_CHARGES: mulliken_charges}

def get_hirsfeld_charges(lines):
    """ Read Hirsfeld charges - run a NBO calculation"""
    keywords = ["Hirshfeld charges,", "Hirshfeld charges with"]
    start, stop = linesio.get_indices_patterns(lines, keywords)
    hirshfeld_charges = [float(line.split()[2]) for line in lines[start + 2 : stop - 1]]

    return {COLUMN_HIRSHFELD_CHARGES: hirshfeld_charges}

def get_cm5_charges(lines):
    """ Read CM5 charges - run a NBO calculation"""
    keywords = ["Hirshfeld charges,", "Hirshfeld charges with"]
    start, stop = linesio.get_indices_patterns(lines, keywords)
    cm5_charges = [float(line.split()[-1]) for line in lines[start + 2 : stop - 1]]

    return {COLUMN_CM5_CHARGES: cm5_charges}

def get_nbo_bond_orders(lines):
    """
    Read Wiberg index - run a NBOread calculation.
    N.B. add $nbo bndidx $end to the footer.
    """
    keywords = ["Wiberg bond index matrix", "Wiberg bond index"]
    start, stop = linesio.get_indices_patterns(lines, keywords)

    if start is None:
        _logger.critical("You did not add: '$nbo bndidx $end' to the footer")
        return None

    # Extract Bond order matrix
    bond_idx_blocks = []
    block_idx = 0
    for line in lines[start + 1 : stop]:
        if line := line.strip().split():  # filter empty strings
            if "Atom" in line:
                continue

            if set("".join(line)).issubset(set("-")):
                bond_idx_blocks.append([])
                block_idx += 1
                continue

            bond_idx_blocks[block_idx - 1].append([float(x) for x in line[2:]])

    # Format matrix
    bond_order_matrix = bond_idx_blocks[0]
    for bond_idx_block in bond_idx_blocks[1:]:
        for i, bond_idx in enumerate(bond_idx_block):
            bond_order_matrix[i].extend(bond_idx)

    return {COLUMN_NBO_BONDORDER: bond_order_matrix}

def get_nmr_shielding_constants(lines):
    """ Read GIAO NMR shielding constants """
    keywords = ["Magnetic shielding tensor (ppm):", "************"]
    start, stop = linesio.get_rev_indices_patterns(lines, keywords)

    shielding_constants = []
    for line in lines[start:stop]:
        if "Isotropic" in line:
            shielding_constants.append(float(line.strip().split()[4]))

    return {COLUMN_SHIELDING_CONSTANTS: shielding_constants}
<commit_msg>Remove input print statement<commit_after>import logging
import pathlib
import tempfile
from collections import ChainMap

from tqdm import tqdm

from ppqm import chembridge, constants, linesio, shell, units
from ppqm.calculator import BaseCalculator

G16_CMD = "g16"
G16_FILENAME = "_tmp_g16_input.com"

_logger = logging.getLogger("g16")

COLUMN_SCF_ENERGY = "scf_energy"
COLUMN_CONVERGED = "is_converged"
COLUMN_MULIKEN_CHARGES = "mulliken charges"
COLUMN_CM5_CHARGES = "cm5_charges"
COLUMN_HIRSHFELD_CHARGES = "hirshfeld_charges"
COLUMN_NBO_BONDORDER = "bond_orders"
COLUMN_SHIELDING_CONSTANTS = "shielding_constants"

class GaussianCalculator(BaseCalculator):
    def __init__(
        self,
        cmd=G16_CMD,
        filename=G16_FILENAME,
        show_progress=False,
        n_cores=None,
        memory=2,
        **kwargs,
    ):

        super().__init__(**kwargs)

        self.cmd = cmd
        self.filename = filename
        self.n_cores = n_cores
        self.memory = memory
        self.show_progress = show_progress

        self.g16_options = {
            "cmd": self.cmd,
            "scr": self.scr,
            "filename": self.filename,
            "memory": self.memory,
        }

        # Default G16 options
        self.options = {}

        #
        self.health_check()

    def __repr__(self) -> str:
        return f"G16Calc(cmd={self.cmd}, scr={self.scr}, n_cores={self.n_cores}, memory={self.memory}gb)"

    def health_check(self):
        """ """

    def calculate(self, molobj, options, footer=None):
        """ """

        if self.n_cores > 1:
            raise NotImplementedError("Parallel not implemented yet.")
        else:
            results = self.calculate_serial(molobj, options, footer=footer)

        return results

    def calculate_serial(self, molobj, options, footer=None):
        """ """

        # If not singlet "spin" is part of options
        if "spin" in options.keys():
            spin = str(options.pop("spin"))
        else:
            spin = str(1)

        options_prime = ChainMap(options, self.options)
        options_prime = dict(options_prime)

        n_confs = molobj.GetNumConformers()
        atoms, _, charge = chembridge.get_axyzc(molobj, atomfmt=str)

        if self.show_progress:
            pbar = tqdm(
                total=n_confs,
                desc="g16(1)",
                **constants.TQDM_OPTIONS,
            )

        properties_list = []
        for conf_idx in range(n_confs):

            coord = chembridge.get_coordinates(molobj, confid=conf_idx)

            properties = get_properties_from_axyzc(
                atoms,
                coord,
                charge,
                spin,
                options=options_prime,
                footer=footer,
                **self.g16_options,
            )

            properties_list.append(properties)

            if self.show_progress:
                pbar.update(1)

        if self.show_progress:
            pbar.close()

        return properties_list

def get_properties_from_axyzc(
    atoms_str,
    coordinates,
    charge,
    spin,
    options=None,
    scr=constants.SCR,
    clean_tempfiles=False,
    cmd=G16_CMD,
    filename=G16_FILENAME,
    footer=None,
    **kwargs,
):

    if isinstance(scr, str):
        scr = pathlib.Path(scr)

    if not filename.endswith(".com"):
        filename += ".com"

    tempdir = tempfile.TemporaryDirectory(dir=scr, prefix="g16_")
    scr = pathlib.Path(tempdir.name)

    # write input file
    input_header = get_header(options, **kwargs)

    inputstr = get_inputfile(
        atoms_str, coordinates, charge, spin, input_header, footer=footer, title="g16 input"
    )

    with open(scr / filename, "w") as f:
        f.write(inputstr)

    # Run subprocess cmd
    cmd = " ".join([cmd, filename])
    _logger.debug(cmd)

    lines = shell.stream(cmd, cwd=scr)
    lines = list(lines)

    # for line in lines:
    #    print(line.strip())
    # print("="*80 )
    # print()

    termination_pattern = "Normal termination of Gaussian"
    idx = linesio.get_rev_index(lines, termination_pattern, stoppattern="File lengths")
    if idx is None:
        _logger.critical("Abnormal termination of Gaussian")
        return None

    # Parse properties from Gaussian output
    properties = read_properties(lines, options)

    # Clean scr dir. TODO: Does this work??
    if clean_tempfiles:
        tempdir.cleanup()

    return properties

def get_inputfile(atom_strs, coordinates, charge, spin, header, footer=None, **kwargs):
    """ """

    inputstr = header + 2 * "\n"
    inputstr += f"  {kwargs['title']}" + 2 * "\n"

    inputstr += f"{charge}  {spin} \n"
    for atom_str, coord in zip(atom_strs, coordinates):
        inputstr += f"{atom_str}  " + " ".join([str(x) for x in coord]) + "\n"
    inputstr += "\n"  # magic line

    if footer is not None:
        inputstr += footer + "\n"

    return inputstr

def get_header(options, **kwargs):
    """ Write G16 header """

    header = f"%mem={kwargs.pop('memory')}gb\n"
    header += "# "
    for key, value in options.items():
        if (value is None) or (len(value) == 0):
            header += f"{key} "
        elif isinstance(value, str):
            header += f"{key}={value} "
        else:
            header += f"{key}=("
            for subkey, subvalue in value.items():
                if (subvalue is None) or (len(subvalue) == 0):
                    header += f"{subkey}, "
                else:
                    header += f"{subkey}={subvalue}, "
            header = header[:-2] + ") "

    return header

def read_properties(lines, options):
    """ Extract values from output depending on calculation options """

    # Collect readers
    readers = []

    if "opt" in options:
        raise NotImplementedError("not implemented opt properties parser")
        # reader = read_properties_opt
    else:
        readers.append(read_properties_sp)

    if "pop" in options:
        if "nboread" in options["pop"]:
            readers.append(get_nbo_bond_orders)

        if "hirsfeld" in options["pop"]:
            readers.append(get_hirsfeld_charges)
            readers.append(get_cm5_charges)

    if "nmr" in options:
        readers.append(get_nmr_shielding_constants)

    # Get properties
    properties = dict()
    for reader in readers:
        new_properties = reader(lines)
        assert isinstance(new_properties, dict)

        properties.update(new_properties)

    return properties

def read_properties_sp(lines):
    """
    Read Mulliken charges
    """
    properties = dict()

    for line in lines:
        if "SCF Done:  E(" in line:
            scf_energy = float(line.split()[4]) * units.hartree_to_kcalmol
            properties[COLUMN_SCF_ENERGY] = scf_energy
            break

    properties.update(get_mulliken_charges(lines))

    return properties

def read_properties_opt(lines):
    """ """

def get_mulliken_charges(lines):
    """ Read Mulliken charges """
    keywords = ["Mulliken charges:", "Sum of Mulliken charges"]
    start, stop = linesio.get_rev_indices_patterns(lines, keywords)
    mulliken_charges = [float(x.split()[-1]) for x in lines[start + 2 : stop]]
    return {COLUMN_MULIKEN_CHARGES: mulliken_charges}

def get_hirsfeld_charges(lines):
    """ Read Hirsfeld charges - run a NBO calculation"""
    keywords = ["Hirshfeld charges,", "Hirshfeld charges with"]
    start, stop = linesio.get_indices_patterns(lines, keywords)
    hirshfeld_charges = [float(line.split()[2]) for line in lines[start + 2 : stop - 1]]

    return {COLUMN_HIRSHFELD_CHARGES: hirshfeld_charges}

def get_cm5_charges(lines):
    """ Read CM5 charges - run a NBO calculation"""
    keywords = ["Hirshfeld charges,", "Hirshfeld charges with"]
    start, stop = linesio.get_indices_patterns(lines, keywords)
    cm5_charges = [float(line.split()[-1]) for line in lines[start + 2 : stop - 1]]

    return {COLUMN_CM5_CHARGES: cm5_charges}

def get_nbo_bond_orders(lines):
    """
    Read Wiberg index - run a NBOread calculation.
    N.B. add $nbo bndidx $end to the footer.
    """
    keywords = ["Wiberg bond index matrix", "Wiberg bond index"]
    start, stop = linesio.get_indices_patterns(lines, keywords)

    if start is None:
        _logger.critical("You did not add: '$nbo bndidx $end' to the footer")
        return None

    # Extract Bond order matrix
    bond_idx_blocks = []
    block_idx = 0
    for line in lines[start + 1 : stop]:
        if line := line.strip().split():  # filter empty strings
            if "Atom" in line:
                continue

            if set("".join(line)).issubset(set("-")):
                bond_idx_blocks.append([])
                block_idx += 1
                continue

            bond_idx_blocks[block_idx - 1].append([float(x) for x in line[2:]])

    # Format matrix
    bond_order_matrix = bond_idx_blocks[0]
    for bond_idx_block in bond_idx_blocks[1:]:
        for i, bond_idx in enumerate(bond_idx_block):
            bond_order_matrix[i].extend(bond_idx)

    return {COLUMN_NBO_BONDORDER: bond_order_matrix}

def get_nmr_shielding_constants(lines):
    """ Read GIAO NMR shielding constants """
    keywords = ["Magnetic shielding tensor (ppm):", "************"]
    start, stop = linesio.get_rev_indices_patterns(lines, keywords)

    shielding_constants = []
    for line in lines[start:stop]:
        if "Isotropic" in line:
            shielding_constants.append(float(line.strip().split()[4]))

    return {COLUMN_SHIELDING_CONSTANTS: shielding_constants}
<|endoftext|>


<commit_before>from django_simple_response.utils import is_geo_value, is_model_instance, is_valuesset
from django_simple_response.utils import is_related_manager, is_queryset, is_iter, is_generator

def basic_serialization(value):
    return unicode(value)

def careful_serialization(value):
    if is_geo_value(value):
        return value.geojson
    if is_model_instance(value):
        return {f.name: getattr(value, f.name) for f in value._meta.fields}
    if is_valuesset(value):
        return [dict(i) for i in value]
    if is_related_manager(value):
        return [i for i in value.all()]
    if is_iter(value): # or is_queryset or is_generator
        return [i for i in value]
    return unicode(value)
<commit_msg>bugfix basic serialization<commit_after>from django_simple_response.utils import is_geo_value, is_model_instance, is_valuesset
from django_simple_response.utils import is_related_manager, is_queryset, is_iter, is_generator

def basic_serialization(value):
    try:
        return unicode(value)
    except StopIteration:
        return repr(value)

def careful_serialization(value):
    if is_geo_value(value):
        return value.geojson
    if is_model_instance(value):
        return {f.name: getattr(value, f.name) for f in value._meta.fields}
    if is_valuesset(value):
        return [dict(i) for i in value]
    if is_related_manager(value):
        return [i for i in value.all()]
    if is_iter(value): # or is_queryset or is_generator
        return [i for i in value]
    return unicode(value)
<|endoftext|>


<commit_before>from django import forms
from .models import Champion, Position

class DraftForm(forms.Form):
    CHAMPIONS = Champion.objects.order_by('display_name')
    CHAMPION_CHOICES = []
    for champion in CHAMPIONS:
        CHAMPION_CHOICES.append((champion.id, champion.display_name))

    PICK_CHOICES = [("","")]+CHAMPION_CHOICES
    NO_BAN = -1
    BAN_CHOICES = [("",""),(NO_BAN,"No Ban")]+CHAMPION_CHOICES

    POSITION_CHOICES = [(None,""),
                        (1,"ADC"),
                        (2,"MID"),
                        (3,"TOP"),
                        (4,"JNG"),
                        (5,"SUP")]

    blue_style_class = "blue"
    red_style_class = "red"
    # fuck everything below this line
    widget_attrs = {'class':blue_style_class}
    blue_ban_0 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Ban 0", initial="", choices=BAN_CHOICES, required=False)
    blue_ban_1 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Ban 1", initial="", choices=BAN_CHOICES, required=False)
    blue_ban_2 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Ban 2", initial="", choices=BAN_CHOICES, required=False)
    blue_ban_3 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Ban 3", initial="", choices=BAN_CHOICES, required=False)
    blue_ban_4 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Ban 4", initial="", choices=BAN_CHOICES, required=False)

    blue_pick_0 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Pick 0", initial="", choices=PICK_CHOICES, required=False)
    blue_pick_1 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Pick 1", initial="", choices=PICK_CHOICES, required=False)
    blue_pick_2 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Pick 2", initial="", choices=PICK_CHOICES, required=False)
    blue_pick_3 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Pick 3", initial="", choices=PICK_CHOICES, required=False)
    blue_pick_4 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Pick 4", initial="", choices=PICK_CHOICES, required=False)

    blue_pos_0 = forms.TypedChoiceField(coerce=int, empty_value=None, widget=forms.Select(attrs=widget_attrs), label="Pos", choices=POSITION_CHOICES, required=False)
    blue_pos_1 = forms.TypedChoiceField(coerce=int, empty_value=None, widget=forms.Select(attrs=widget_attrs), label="Pos", choices=POSITION_CHOICES, required=False)
    blue_pos_2 = forms.TypedChoiceField(coerce=int, empty_value=None, widget=forms.Select(attrs=widget_attrs), label="Pos", choices=POSITION_CHOICES, required=False)
    blue_pos_3 = forms.TypedChoiceField(coerce=int, empty_value=None, widget=forms.Select(attrs=widget_attrs), label="Pos", choices=POSITION_CHOICES, required=False)
    blue_pos_4 = forms.TypedChoiceField(coerce=int, empty_value=None, widget=forms.Select(attrs=widget_attrs), label="Pos", choices=POSITION_CHOICES, required=False)

    widget_attrs = {'class':red_style_class}
    red_ban_0 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Ban 0", initial="", choices=BAN_CHOICES, required=False)
    red_ban_1 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Ban 1", initial="", choices=BAN_CHOICES, required=False)
    red_ban_2 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Ban 2", initial="", choices=BAN_CHOICES, required=False)
    red_ban_3 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Ban 3", initial="", choices=BAN_CHOICES, required=False)
    red_ban_4 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Ban 4", initial="", choices=BAN_CHOICES, required=False)

    red_pick_0 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Pick 0", initial="", choices=PICK_CHOICES, required=False)
    red_pick_1 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Pick 1", initial="", choices=PICK_CHOICES, required=False)
    red_pick_2 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Pick 2", initial="", choices=PICK_CHOICES, required=False)
    red_pick_3 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Pick 3", initial="", choices=PICK_CHOICES, required=False)
    red_pick_4 = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Pick 4", initial="", choices=PICK_CHOICES, required=False)

    red_pos_0 = forms.TypedChoiceField(coerce=int, empty_value=None, widget=forms.Select(attrs=widget_attrs), label="Pos", choices=POSITION_CHOICES, required=False)
    red_pos_1 = forms.TypedChoiceField(coerce=int, empty_value=None, widget=forms.Select(attrs=widget_attrs), label="Pos", choices=POSITION_CHOICES, required=False)
    red_pos_2 = forms.TypedChoiceField(coerce=int, empty_value=None, widget=forms.Select(attrs=widget_attrs), label="Pos", choices=POSITION_CHOICES, required=False)
    red_pos_3 = forms.TypedChoiceField(coerce=int, empty_value=None, widget=forms.Select(attrs=widget_attrs), label="Pos", choices=POSITION_CHOICES, required=False)
    red_pos_4 = forms.TypedChoiceField(coerce=int, empty_value=None, widget=forms.Select(attrs=widget_attrs), label="Pos", choices=POSITION_CHOICES, required=False)
<commit_msg>forms now dynamically generate<commit_after>from django import forms
from .models import Champion, Position

class DraftForm(forms.Form):
    def __init__(self, *args, **kwargs):
        super(DraftForm, self).__init__(*args, **kwargs)
        teams = ["blue", "red"]
        NUM_PICKS = 5
        NUM_BANS = 5

        self.POSITION_CHOICES = [(None,""),
                            (1,"ADC"),
                            (2,"MID"),
                            (3,"TOP"),
                            (4,"JNG"),
                            (5,"SUP")]

        CHAMPIONS = Champion.objects.order_by('display_name')

        CHAMPION_CHOICES = []
        for champion in CHAMPIONS:
            CHAMPION_CHOICES.append((champion.id, champion.display_name))

        PICK_CHOICES = [("","")]+CHAMPION_CHOICES
        NO_BAN = -1
        BAN_CHOICES = [("",""),(NO_BAN,"No Ban")]+CHAMPION_CHOICES

        blue_style_class = "blue"
        red_style_class = "red"
        styles = {"blue":blue_style_class, "red":red_style_class}

        for team in teams:
            widget_attrs = {"class": styles[team]}
            for ban_count in range(NUM_BANS):
                name = "{}_{}_{}".format(team, "ban", ban_count)
                self.fields[name] = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Ban {}".format(ban_count), initial="", choices=BAN_CHOICES, required=False)

            for pick_count in range(NUM_PICKS):
                name = "{}_{}_{}".format(team, "pick", pick_count)
                self.fields[name] = forms.ChoiceField(widget=forms.Select(attrs=widget_attrs), label="Pick {}".format(pick_count), initial="", choices=PICK_CHOICES, required=False)

                name = "{}_{}_{}".format(team, "pos", pick_count)
                self.fields[name] = forms.TypedChoiceField(coerce=int, empty_value=None, widget=forms.Select(attrs=widget_attrs), label="Pos", choices=self.POSITION_CHOICES, required=False)
<|endoftext|>


<commit_before>import sys
import cmd

from django.core.management.base import BaseCommand
from django.utils.text import smart_split

from djapian import utils
from djapian import IndexSpace

def with_index(func):
    def _decorator(cmd, arg):
        if cmd._current_index is None:
            print "No index selected"
            return

        return func(cmd, arg)
    _decorator.__doc__ = func.__doc__
    return _decorator

def split_arg(func):
    def _decorator(cmd, arg):
        bits = list(smart_split(arg))

        return func(cmd, *bits)
    _decorator.__doc__ = func.__doc__
    return _decorator

class Interpreter(cmd.Cmd):
    prompt = ">>> "

    def __init__(self, *args):
        self._current_index = None

        if len(args):
            self.do_use(args[0])

        cmd.Cmd.__init__(self)

    def do_list(self, arg):
        """
        Lists all available indexes with their ids
        """
        print "Installed spaces/models/indexers:"
        for space_i, space in enumerate(IndexSpace.instances):
            print "%s: `%s`" % (space_i, space)
            for model_indexer_i, pair in enumerate(space.get_indexers().items()):
                model, indexers = pair
                print "  %s.%s: `%s`" % (space_i, model_indexer_i, utils.model_name(model))
                for indexer_i, indexer in enumerate(indexers):
                    print "    %s.%s.%s: `%s`" % (space_i, model_indexer_i, indexer_i, indexer)

    def do_exit(self, arg):
        """
        Exit shell
        """
        return True

    def do_use(self, index):
        """
        Changes current index
        """
        space, model, indexer = self._get_indexer(index)

        if indexer is not None:
            self._current_index = indexer

            print "Using `%s:%s:%s` index." % (space, utils.model_name(model), indexer)

    def do_usecomposite(self, indexes):
        """
        Changes current index to composition of given indexers
        """
        from djapian.indexer import CompositeIndexer

        indexers = []
        for index in indexes.split(' '):
            indexers.append(self._get_indexer(index.strip()))

        self._current_index = CompositeIndexer(*[i[2] for i in indexers])

        print "Using composition of:"
        for indexer in indexers:
            space, model, indexer = indexer
            print "  `%s:%s:%s`" % (space, utils.model_name(model), indexer)
        print "indexes."

    @with_index
    @split_arg
    def do_query(self, query, _slice=''):
        """
        Returns objects fetched by given query
        """
        _slice = slice(*self._parse_slice(_slice))
        print list(self._current_index.search(query)[_slice])

    @with_index
    def do_count(self, query):
        """
        Returns count of objects fetched by given query
        """
        print self._current_index.search(query).count()

    @with_index
    def do_total(self, arg):
        """
        Returns count of objects in index
        """
        print self._current_index.document_count()

    def do_stats(self, arg):
        """
        Print index status information
        """
        import operator
        print "Number of spaces: %s" % len(IndexSpace.instances)
        print "Number of indexes: %s" % reduce(
            operator.add,
            [len(space.get_indexers()) for space in IndexSpace.instances]
        )

    @with_index
    def do_docslist(self, slice=""):
        """
        Returns count of objects in index
        """
        db = self._current_index._db.open()

        start, end = self._parse_slice(slice, default=(1, db.get_lastdocid()))

        for i in range(start, end + 1):
            doc = db.get_document(i)
            print "doc #%s:\n\tValues (%s):" % (i, doc.values_count())
            val = doc.values_begin()

            for i in range(doc.values_count()):
                print "\t\t%s: %s" % (val.get_valueno(), val.get_value())
                val.next()

            print "\tTerms (%s):" % doc.termlist_count()
            termlist = doc.termlist_begin()

            for i in range(doc.termlist_count()):
                print termlist.get_term(),
                termlist.next()
            print "\n"

    @with_index
    def do_delete(self, id):
        """
        Removes document by id
        """
        id = int(id)

        db = self._current_index._db.open(write=True)
        db.delete_document(id)
        print "Document #%s deleted." % id

    def _get_indexer(self, index):
        try:
            space, model, indexer = self._parse_slice(index, '.')

            space = IndexSpace.instances[space]
            model = space.get_indexers().keys()[model]
            indexer = space.get_indexers()[model][indexer]
        except (TypeError, IndexError, KeyError, ValueError):
            print 'Illegal index alias `%s`. See `list` command for available aliases' % index
            return None, None, None

        return space, model, indexer

    def _parse_slice(self, slice='', delimeter=':', default=tuple()):
        if slice:
            def _select(b, d):
                try:
                    return int(b)
                except ValueError:
                    return d

            bits = [_select(b, d) for b, d in zip(slice.split(delimeter), default or ([None] * 3))]
        elif default:
            return default
        else:
            raise ValueError("Empty slice")

        return bits

class Command(BaseCommand):
    help = "Djapian shell that provides capabilities to monitoring indexes."
    args = '[index_id]'

    requires_model_validation = True

    def handle(self, *args, **options):
        utils.load_indexes()

        try:
            Interpreter(*args).cmdloop("Interactive Djapian shell.")
        except KeyboardInterrupt:
            print "\n"
<commit_msg>Added current index highlightning in indexshell `list` command<commit_after>import sys
import cmd

from django.core.management.base import BaseCommand
from django.utils.text import smart_split

from djapian import utils
from djapian import IndexSpace

def with_index(func):
    def _decorator(cmd, arg):
        if cmd._current_index is None:
            print "No index selected"
            return

        return func(cmd, arg)
    _decorator.__doc__ = func.__doc__
    return _decorator

def split_arg(func):
    def _decorator(cmd, arg):
        bits = list(smart_split(arg))

        return func(cmd, *bits)
    _decorator.__doc__ = func.__doc__
    return _decorator

class Interpreter(cmd.Cmd):
    prompt = ">>> "

    def __init__(self, *args):
        self._current_index = None
        self._current_index_path = [None, None, None]

        if len(args):
            self.do_use(args[0])

        cmd.Cmd.__init__(self)

    def do_list(self, arg):
        """
        Lists all available indexes with their ids
        """
        print "Installed spaces/models/indexers:"

        def _is_selected(space, model=None, index=None):
            selected = [b for b in (space, model, index) if b is not None]
            return len(
                [1 for c, s in zip(self._current_index_path, selected) if c == s]
            ) == len(selected)

        for space_i, space in enumerate(IndexSpace.instances):
            print (_is_selected(space_i) and '* ' or '- ') + '%s: `%s`' % (space_i, space)
            for model_indexer_i, pair in enumerate(space.get_indexers().items()):
                model, indexers = pair
                print (_is_selected(space_i, model_indexer_i) and '  * ' or '  - ') +\
                        "%s.%s: `%s`" % (space_i, model_indexer_i, utils.model_name(model))
                for indexer_i, indexer in enumerate(indexers):
                    print (_is_selected(space_i, model_indexer_i, indexer_i) and '    * ' or '    - ') +\
                        "%s.%s.%s: `%s`" % (space_i, model_indexer_i, indexer_i, indexer)

    def do_exit(self, arg):
        """
        Exit shell
        """
        return True

    def do_use(self, index):
        """
        Changes current index
        """
        space, model, indexer, path = self._get_indexer(index)

        if indexer is not None:
            self._current_index = indexer
            self._current_index_path = path

            print "Using `%s:%s:%s` index." % (space, utils.model_name(model), indexer)

    def do_usecomposite(self, indexes):
        """
        Changes current index to composition of given indexers
        """
        from djapian.indexer import CompositeIndexer

        indexers = []
        for index in indexes.split(' '):
            indexers.append(self._get_indexer(index.strip()))

        self._current_index = CompositeIndexer(*[i[2] for i in indexers])

        print "Using composition of:"
        for indexer in indexers:
            space, model, indexer = indexer
            print "  `%s:%s:%s`" % (space, utils.model_name(model), indexer)
        print "indexes."

    @with_index
    @split_arg
    def do_query(self, query, _slice=''):
        """
        Returns objects fetched by given query
        """
        _slice = slice(*self._parse_slice(_slice))
        print list(self._current_index.search(query)[_slice])

    @with_index
    def do_count(self, query):
        """
        Returns count of objects fetched by given query
        """
        print self._current_index.search(query).count()

    @with_index
    def do_total(self, arg):
        """
        Returns count of objects in index
        """
        print self._current_index.document_count()

    def do_stats(self, arg):
        """
        Print index status information
        """
        import operator
        print "Number of spaces: %s" % len(IndexSpace.instances)
        print "Number of indexes: %s" % reduce(
            operator.add,
            [len(space.get_indexers()) for space in IndexSpace.instances]
        )

    @with_index
    def do_docslist(self, slice=""):
        """
        Returns count of objects in index
        """
        db = self._current_index._db.open()

        start, end = self._parse_slice(slice, default=(1, db.get_lastdocid()))

        for i in range(start, end + 1):
            doc = db.get_document(i)
            print "doc #%s:\n\tValues (%s):" % (i, doc.values_count())
            val = doc.values_begin()

            for i in range(doc.values_count()):
                print "\t\t%s: %s" % (val.get_valueno(), val.get_value())
                val.next()

            print "\tTerms (%s):" % doc.termlist_count()
            termlist = doc.termlist_begin()

            for i in range(doc.termlist_count()):
                print termlist.get_term(),
                termlist.next()
            print "\n"

    @with_index
    def do_delete(self, id):
        """
        Removes document by id
        """
        id = int(id)

        db = self._current_index._db.open(write=True)
        db.delete_document(id)
        print "Document #%s deleted." % id

    def _get_indexer(self, index):
        try:
            _space, _model, _indexer = self._parse_slice(index, '.')

            space = IndexSpace.instances[_space]
            model = space.get_indexers().keys()[_model]
            indexer = space.get_indexers()[model][_indexer]
        except (TypeError, IndexError, KeyError, ValueError):
            print 'Illegal index alias `%s`. See `list` command for available aliases' % index
            return None, None, None, None

        return space, model, indexer, [_space, _model, _indexer]

    def _parse_slice(self, slice='', delimeter=':', default=tuple()):
        if slice:
            def _select(b, d):
                try:
                    return int(b)
                except ValueError:
                    return d

            bits = [_select(b, d) for b, d in zip(slice.split(delimeter), default or ([None] * 3))]
        elif default:
            return default
        else:
            raise ValueError("Empty slice")

        return bits

class Command(BaseCommand):
    help = "Djapian shell that provides capabilities to monitoring indexes."
    args = '[index_id]'

    requires_model_validation = True

    def handle(self, *args, **options):
        utils.load_indexes()

        try:
            Interpreter(*args).cmdloop("Interactive Djapian shell.")
        except KeyboardInterrupt:
            print "\n"
<|endoftext|>