<commit_before>#!/usr/bin/env python
"""
make_a_star_cluster.py creates a model star cluster,
which can then be used in N-body simulations or for other purposes.

It requires AMUSE, which can be downloaded from http://amusecode.org or
https://github.com/amusecode/amuse.

Currently not feature-complete yet, and function/argument names are
subject to change.

-- Steven Rieder steven at rieder punt nl
"""
from __future__ import print_function, division

import os
import argparse

import numpy as np

# from amuse.lab import *
from amuse.io import write_set_to_file
from amuse.lab import (
    units,
    nbody_system,
    generic_unit_converter,
)
from amuse.support.console import set_printing_strategy

def new_argument_parser():
    "Parse command line arguments"
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-o',
        dest='clustername',
        default="auto",
        help="Outputfile name [auto]",
    )
    parser.add_argument(
        '-type',
        dest='filetype',
        default="amuse",
        help="Output file type ([amuse]/txt/starlab/nemo)",
        )
    parser.add_argument(
        '-N',
        dest='number_of_stars',
        default=1024,
        type=int,
        help='Number of stars [1024]',
    )
    parser.add_argument(
        '-M',
        dest='cluster_mass',
        type=float,
        default=0,
        help='Cluster mass (takes precedence over number of stars)',
    )
    parser.add_argument(
        '-dist',
        dest='star_distribution',
        default="plummer",
        help="Star distribution ([plummer]/king/fractal)",
    )
    parser.add_argument(
        '-gasdist',
        dest='gas_distribution',
        default="none",
        help="Gas distribution ([none]/plummer/king/fractal) \
                NOT IMPLEMENTED YET",
    )
    parser.add_argument(
        '-imf',
        dest='initial_mass_function',
        default="kroupa",
        help="IMF ([kroupa]/salpeter/fixed)",
    )
    parser.add_argument(
        '-modelnr',
        dest='cluster_model_number',
        type=int,
        default=30483,
        help='Model number',
    )
    parser.add_argument(
        '-kingw',
        dest='king_parameter_w0',
        type=float,
        default=7.0,
        help='King W0 parameter [7.0]',
    )
    parser.add_argument(
        '-fractald',
        dest='fractal_parameter_fd',
        type=float,
        default=2.0,
        help='Fractal distribution [2.0]',
    )
    parser.add_argument(
        '-uml',
        dest='upper_mass_limit',
        type=float,
        default=125.,
        help='Upper mass limit (in MSun) [125.]',
    )
    parser.add_argument(
        '-Z',
        dest='metallicity',
        type=float,
        default=0.01,
        help='Metallicity [0.01]',
    )
    parser.add_argument(
        '-ibf',
        dest='initial_binary_fraction',
        type=float,
        default=0.,
        help='Initial binary fraction (0-1) [0.] NOT IMPLEMENTED YET',
    )
    parser.add_argument(
        '-R',
        dest='effective_radius',
        type=float,
        default=3.0,
        help='Cluster radius [3.0] (parsec)',
    )
    parser.add_argument(
        '-vr',
        dest='virial_ratio',
        type=float,
        default=0.5,
        help="Virial ration [0.5], 0.5=stable, 0.75=just expelled gas, \
                0.1=collapsing",
    )
    args = parser.parse_args()
    return args

def make_a_star_cluster(
        stellar_mass=False,
        initial_mass_function="salpeter",
        upper_mass_limit=125. | units.MSun,
        number_of_stars=1024,
        effective_radius=3.0 | units.parsec,
        star_distribution="plummer",
        star_distribution_w0=7.0,
        star_distribution_fd=2.0,
        star_metallicity=0.01,
        # initial_binary_fraction=0,
        **kwargs
):
    """
    Create stars.
    When using an IMF, either the stellar mass is fixed (within
    stochastic error) or the number of stars is fixed. When using
    equal-mass stars, both are fixed.
    """

    if stellar_mass:
        # Add stars to cluster, until mass limit reached (inclusive!)
        if initial_mass_function == "kroupa":
            from amuse.ic.brokenimf import new_kroupa_mass_distribution
            mass = new_kroupa_mass_distribution(0)
            while mass.sum() < stellar_mass:
                mass.append(
                    new_kroupa_mass_distribution(
                        1,
                        mass_max=upper_mass_limit,
                        )[0]
                )
            total_mass = mass.sum()
            number_of_stars = len(mass)
        elif initial_mass_function == "salpeter":
            from amuse.ic.salpeter import new_salpeter_mass_distribution
            mass = new_salpeter_mass_distribution(0)
            while mass.sum() < stellar_mass:
                mass.append(
                    new_salpeter_mass_distribution(
                        1,
                        mass_max=upper_mass_limit,
                        )[0]
                )
            total_mass = mass.sum()
            number_of_stars = len(mass)
        elif initial_mass_function == "fixed":
            mass_of_individual_star = stellar_mass / number_of_stars
            mass = mass_of_individual_star
            total_mass = stellar_mass
        else:
            return -1, "No mass function"
    else:
        # Give stars their mass
        if initial_mass_function == "kroupa":
            from amuse.ic.brokenimf import new_kroupa_mass_distribution
            mass = new_kroupa_mass_distribution(
                number_of_stars,
                mass_max=upper_mass_limit,
            )
            total_mass = mass.sum()
        elif initial_mass_function == "salpeter":
            from amuse.ic.salpeter import new_salpeter_mass_distribution
            mass = new_salpeter_mass_distribution(
                number_of_stars,
                mass_max=upper_mass_limit,
            )
            total_mass = mass.sum()
        elif initial_mass_function == "fixed":
            mass = mass_of_individual_star
            total_mass = number_of_stars * mass
        else:
            return -1, "No mass function"

    converter = generic_unit_converter.ConvertBetweenGenericAndSiUnits(
        total_mass,
        1. | units.kms,
        effective_radius,
    )
    """
    Give stars a position and velocity, based on the distribution
    model.
    """
    if star_distribution == "plummer":
        from amuse.ic.plummer import new_plummer_sphere
        stars = new_plummer_sphere(
            number_of_stars,
            convert_nbody=converter,
        )
    elif star_distribution == "king":
        from amuse.ic.kingmodel import new_king_model
        stars = new_king_model(
            number_of_stars,
            star_distribution_w0,
            convert_nbody=converter,
        )
    elif star_distribution == "fractal":
        from amuse.ic.fractalcluster import new_fractal_cluster_model
        stars = new_fractal_cluster_model(
            number_of_stars,
            fractal_dimension=star_distribution_fd,
            convert_nbody=converter,
        )
    else:
        return -1, "No stellar distribution"

    """
    set the stellar mass.
    """
    stars.mass = mass

    """
    set other stellar parameters.
    """
    stars.metallicity = star_metallicity

    """
    Virialize the star cluster
    """
    stars.move_to_center()
    stars.scale_to_standard(
        convert_nbody=converter,
        # virial_ratio=virial_ratio,
        # smoothing_length_squared= ...,
    )

    """
    Record the cluster's initial parameters to the particle distribution
    """
    stars.collection_attributes.initial_mass_function = initial_mass_function
    stars.collection_attributes.upper_mass_limit = upper_mass_limit
    stars.collection_attributes.number_of_stars = number_of_stars

    stars.collection_attributes.effective_radius = effective_radius

    stars.collection_attributes.star_distribution = star_distribution
    stars.collection_attributes.star_distribution_w0 = star_distribution_w0
    stars.collection_attributes.star_distribution_fd = star_distribution_fd

    stars.collection_attributes.star_metallicity = star_metallicity

    # Derived/legacy values
    stars.collection_attributes.converter_mass = \
        converter.to_si(1 | nbody_system.mass)
    stars.collection_attributes.converter_length =\
        converter.to_si(1 | nbody_system.length)
    stars.collection_attributes.converter_speed =\
        converter.to_si(1 | nbody_system.speed)

    return stars

def main():
    "Run make_a_star_cluster"
    set_printing_strategy(
        "custom",
        preferred_units=[units.MSun, units.parsec, units.yr, units.kms],
        precision=5,
    )
    clustertemplate = "TESTCluster_%08i"

    args = new_argument_parser()
    cluster_model_number = args.cluster_model_number
    star_distribution = args.star_distribution
    # gas_distribution = args.gas_distribution
    king_parameter_w0 = args.king_parameter_w0
    fractal_parameter_fd = args.fractal_parameter_fd
    initial_mass_function = args.initial_mass_function
    number_of_stars = args.number_of_stars
    if args.cluster_mass != 0:
        cluster_mass = args.cluster_mass | units.MSun
    else:
        cluster_mass = False
    upper_mass_limit = args.upper_mass_limit | units.MSun
    effective_radius = args.effective_radius | units.parsec
    metallicity = args.metallicity
    # virial_ratio = args.virial_ratio
    filetype = args.filetype

    # not implemented yet
    # initial_binary_fraction = args.initial_binary_fraction

    np.random.seed(cluster_model_number)

    if not (number_of_stars or cluster_mass):
        print("no number of stars or cluster mass given, exiting")
        exit()

    stars = make_a_star_cluster(
        stellar_mass=cluster_mass,
        initial_mass_function=initial_mass_function,
        upper_mass_limit=upper_mass_limit,
        number_of_stars=number_of_stars,
        effective_radius=effective_radius,
        star_distribution=star_distribution,
        star_distribution_w0=king_parameter_w0,
        star_distribution_fd=fractal_parameter_fd,
        star_metallicity=metallicity,
    )

    print(stars.mass.sum())
    print(len(stars))

    if args.clustername != "auto":
        clustername = args.clustername
    else:
        cluster_file_exists = True
        N = -1
        while cluster_file_exists:
            N += 1
            clustername = (
                clustertemplate % N
                + "." + filetype
            )
            cluster_file_exists = os.path.isfile(clustername)

    write_set_to_file(stars, clustername, filetype)

if __name__ in ["__main__"]:
    main()
<commit_msg>Cleanup comments<commit_after>#!/usr/bin/env python
"""
make_a_star_cluster.py creates a model star cluster,
which can then be used in N-body simulations or for other purposes.

It requires AMUSE, which can be downloaded from http://amusecode.org or
https://github.com/amusecode/amuse.

Currently not feature-complete yet, and function/argument names are
subject to change.

-- Steven Rieder steven at rieder punt nl
"""
from __future__ import print_function, division

import os
import argparse

import numpy as np

# from amuse.lab import *
from amuse.io import write_set_to_file
from amuse.lab import (
    units,
    nbody_system,
    generic_unit_converter,
)
from amuse.support.console import set_printing_strategy

def new_argument_parser():
    "Parse command line arguments"
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-o',
        dest='clustername',
        default="auto",
        help="Outputfile name [auto]",
    )
    parser.add_argument(
        '-type',
        dest='filetype',
        default="amuse",
        help="Output file type ([amuse]/txt/starlab/nemo)",
        )
    parser.add_argument(
        '-N',
        dest='number_of_stars',
        default=1024,
        type=int,
        help='Number of stars [1024]',
    )
    parser.add_argument(
        '-M',
        dest='cluster_mass',
        type=float,
        default=0,
        help='Cluster mass (takes precedence over number of stars)',
    )
    parser.add_argument(
        '-dist',
        dest='star_distribution',
        default="plummer",
        help="Star distribution ([plummer]/king/fractal)",
    )
    parser.add_argument(
        '-gasdist',
        dest='gas_distribution',
        default="none",
        help="Gas distribution ([none]/plummer/king/fractal) \
                NOT IMPLEMENTED YET",
    )
    parser.add_argument(
        '-imf',
        dest='initial_mass_function',
        default="kroupa",
        help="IMF ([kroupa]/salpeter/fixed)",
    )
    parser.add_argument(
        '-modelnr',
        dest='cluster_model_number',
        type=int,
        default=30483,
        help='Model number',
    )
    parser.add_argument(
        '-kingw',
        dest='king_parameter_w0',
        type=float,
        default=7.0,
        help='King W0 parameter [7.0]',
    )
    parser.add_argument(
        '-fractald',
        dest='fractal_parameter_fd',
        type=float,
        default=2.0,
        help='Fractal distribution [2.0]',
    )
    parser.add_argument(
        '-uml',
        dest='upper_mass_limit',
        type=float,
        default=125.,
        help='Upper mass limit (in MSun) [125.]',
    )
    parser.add_argument(
        '-Z',
        dest='metallicity',
        type=float,
        default=0.01,
        help='Metallicity [0.01]',
    )
    parser.add_argument(
        '-ibf',
        dest='initial_binary_fraction',
        type=float,
        default=0.,
        help='Initial binary fraction (0-1) [0.] NOT IMPLEMENTED YET',
    )
    parser.add_argument(
        '-R',
        dest='effective_radius',
        type=float,
        default=3.0,
        help='Cluster radius [3.0] (parsec)',
    )
    parser.add_argument(
        '-vr',
        dest='virial_ratio',
        type=float,
        default=0.5,
        help="Virial ration [0.5], 0.5=stable, 0.75=just expelled gas, \
                0.1=collapsing",
    )
    args = parser.parse_args()
    return args

def make_a_star_cluster(
        stellar_mass=False,
        initial_mass_function="salpeter",
        upper_mass_limit=125. | units.MSun,
        number_of_stars=1024,
        effective_radius=3.0 | units.parsec,
        star_distribution="plummer",
        star_distribution_w0=7.0,
        star_distribution_fd=2.0,
        star_metallicity=0.01,
        # initial_binary_fraction=0,
        **kwargs
):
    """
    Create stars.
    When using an IMF, either the stellar mass is fixed (within
    stochastic error) or the number of stars is fixed. When using
    equal-mass stars, both are fixed.
    """

    if stellar_mass:
        # Add stars to cluster, until mass limit reached (inclusive!)
        if initial_mass_function == "kroupa":
            from amuse.ic.brokenimf import new_kroupa_mass_distribution
            mass = new_kroupa_mass_distribution(0)
            while mass.sum() < stellar_mass:
                mass.append(
                    new_kroupa_mass_distribution(
                        1,
                        mass_max=upper_mass_limit,
                        )[0]
                )
            total_mass = mass.sum()
            number_of_stars = len(mass)
        elif initial_mass_function == "salpeter":
            from amuse.ic.salpeter import new_salpeter_mass_distribution
            mass = new_salpeter_mass_distribution(0)
            while mass.sum() < stellar_mass:
                mass.append(
                    new_salpeter_mass_distribution(
                        1,
                        mass_max=upper_mass_limit,
                        )[0]
                )
            total_mass = mass.sum()
            number_of_stars = len(mass)
        elif initial_mass_function == "fixed":
            mass_of_individual_star = stellar_mass / number_of_stars
            mass = mass_of_individual_star
            total_mass = stellar_mass
        else:
            return -1, "No mass function"
    else:
        # Give stars their mass
        if initial_mass_function == "kroupa":
            from amuse.ic.brokenimf import new_kroupa_mass_distribution
            mass = new_kroupa_mass_distribution(
                number_of_stars,
                mass_max=upper_mass_limit,
            )
            total_mass = mass.sum()
        elif initial_mass_function == "salpeter":
            from amuse.ic.salpeter import new_salpeter_mass_distribution
            mass = new_salpeter_mass_distribution(
                number_of_stars,
                mass_max=upper_mass_limit,
            )
            total_mass = mass.sum()
        elif initial_mass_function == "fixed":
            mass = mass_of_individual_star
            total_mass = number_of_stars * mass
        else:
            return -1, "No mass function"

    converter = generic_unit_converter.ConvertBetweenGenericAndSiUnits(
        total_mass,
        1. | units.kms,
        effective_radius,
    )
    # Give stars a position and velocity, based on the distribution model.
    if star_distribution == "plummer":
        from amuse.ic.plummer import new_plummer_sphere
        stars = new_plummer_sphere(
            number_of_stars,
            convert_nbody=converter,
        )
    elif star_distribution == "king":
        from amuse.ic.kingmodel import new_king_model
        stars = new_king_model(
            number_of_stars,
            star_distribution_w0,
            convert_nbody=converter,
        )
    elif star_distribution == "fractal":
        from amuse.ic.fractalcluster import new_fractal_cluster_model
        stars = new_fractal_cluster_model(
            number_of_stars,
            fractal_dimension=star_distribution_fd,
            convert_nbody=converter,
        )
    else:
        return -1, "No stellar distribution"

    # set the stellar mass.
    stars.mass = mass

    # set other stellar parameters.
    stars.metallicity = star_metallicity

    # Virialize the star cluster
    stars.move_to_center()
    stars.scale_to_standard(
        convert_nbody=converter,
        # virial_ratio=virial_ratio,
        # smoothing_length_squared= ...,
    )

    # Record the cluster's initial parameters to the particle distribution
    stars.collection_attributes.initial_mass_function = initial_mass_function
    stars.collection_attributes.upper_mass_limit = upper_mass_limit
    stars.collection_attributes.number_of_stars = number_of_stars

    stars.collection_attributes.effective_radius = effective_radius

    stars.collection_attributes.star_distribution = star_distribution
    stars.collection_attributes.star_distribution_w0 = star_distribution_w0
    stars.collection_attributes.star_distribution_fd = star_distribution_fd

    stars.collection_attributes.star_metallicity = star_metallicity

    # Derived/legacy values
    stars.collection_attributes.converter_mass = \
        converter.to_si(1 | nbody_system.mass)
    stars.collection_attributes.converter_length =\
        converter.to_si(1 | nbody_system.length)
    stars.collection_attributes.converter_speed =\
        converter.to_si(1 | nbody_system.speed)

    return stars

def main():
    "Run make_a_star_cluster"
    set_printing_strategy(
        "custom",
        preferred_units=[units.MSun, units.parsec, units.yr, units.kms],
        precision=5,
    )
    clustertemplate = "TESTCluster_%08i"

    args = new_argument_parser()
    cluster_model_number = args.cluster_model_number
    star_distribution = args.star_distribution
    # gas_distribution = args.gas_distribution
    king_parameter_w0 = args.king_parameter_w0
    fractal_parameter_fd = args.fractal_parameter_fd
    initial_mass_function = args.initial_mass_function
    number_of_stars = args.number_of_stars
    if args.cluster_mass != 0:
        cluster_mass = args.cluster_mass | units.MSun
    else:
        cluster_mass = False
    upper_mass_limit = args.upper_mass_limit | units.MSun
    effective_radius = args.effective_radius | units.parsec
    metallicity = args.metallicity
    # virial_ratio = args.virial_ratio
    filetype = args.filetype

    # not implemented yet
    # initial_binary_fraction = args.initial_binary_fraction

    np.random.seed(cluster_model_number)

    if not (number_of_stars or cluster_mass):
        print("no number of stars or cluster mass given, exiting")
        exit()

    stars = make_a_star_cluster(
        stellar_mass=cluster_mass,
        initial_mass_function=initial_mass_function,
        upper_mass_limit=upper_mass_limit,
        number_of_stars=number_of_stars,
        effective_radius=effective_radius,
        star_distribution=star_distribution,
        star_distribution_w0=king_parameter_w0,
        star_distribution_fd=fractal_parameter_fd,
        star_metallicity=metallicity,
    )

    print(stars.mass.sum())
    print(len(stars))

    if args.clustername != "auto":
        clustername = args.clustername
    else:
        cluster_file_exists = True
        N = -1
        while cluster_file_exists:
            N += 1
            clustername = (
                clustertemplate % N
                + "." + filetype
            )
            cluster_file_exists = os.path.isfile(clustername)

    write_set_to_file(stars, clustername, filetype)

if __name__ in ["__main__"]:
    main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Fri Mar 21 15:11:31 2014

@author: ibackus
@editor: dfleming
-Note: indentation is 4 spaces in this file, not a tab!

"""

__version__ = "$Revision: 1 $"
# $Source$

import pynbody
SimArray = pynbody.array.SimArray
import numpy as np
import math
import gc
import os
import AddBinary
import isaac
import calc_velocity
import ICgen_utils
import ICglobal_settings
global_settings = ICglobal_settings.global_settings

def snapshot_gen(ICobj):
    """
    Generates a tipsy snapshot from the initial conditions object ICobj.

    Returns snapshot, param

        snapshot: tipsy snapshot
        param: dictionary containing info for a .param file
    Note: Code has been edited (dflemin3) such that now it returns a snapshot for a circumbinary disk
    where initial conditions generated assuming star at origin of mass M.  After gas initialized, replaced
    star at origin with binary system who's center of mass lies at the origin and who's mass m1 +m2 = M
    """

    print 'Generating snapshot...'
    # Constants
    G = SimArray(1.0,'G')
    # ------------------------------------
    # Load in things from ICobj
    # ------------------------------------
    print 'Accessing data from ICs'
    settings = ICobj.settings

    # snapshot file name
    snapshotName = settings.filenames.snapshotName
    paramName = settings.filenames.paramName

    # particle positions
    r = ICobj.pos.r
    xyz = ICobj.pos.xyz

    # Number of particles
    nParticles = ICobj.pos.nParticles

    # molecular mass
    m = settings.physical.m

    # star mass
    m_star = settings.physical.M.copy()

    # disk mass
    m_disk = ICobj.sigma.m_disk.copy()
    m_disk = isaac.match_units(m_disk, m_star)[0]

    # mass of the gas particles
    m_particles = m_disk / float(nParticles)

    # re-scale the particles (allows making of low-mass disk)
    m_particles *= settings.snapshot.mScale

    # -------------------------------------------------
    # Assign output
    # -------------------------------------------------
    print 'Assigning data to snapshot'
    # Get units all set up
    m_unit = m_star.units
    pos_unit = r.units

    if xyz.units != r.units:

        xyz.convert_units(pos_unit)

    # time units are sqrt(L^3/GM)
    t_unit = np.sqrt((pos_unit**3)*np.power((G*m_unit), -1)).units
    # velocity units are L/t
    v_unit = (pos_unit/t_unit).ratio('km s**-1')
    # Make it a unit, save value for future conversion
    v_unit_vel = v_unit
    #Ensure v_unit_vel is the same as what I assume it is.
    assert(np.fabs(AddBinary.VEL_UNIT-v_unit_vel)<AddBinary.SMALL),"VEL_UNIT not equal to ChaNGa unit! Why??"

    v_unit = pynbody.units.Unit('{0} km s**-1'.format(v_unit))

    # Other settings
    metals = settings.snapshot.metals
    star_metals = metals

    # Generate snapshot
    # Note that empty pos, vel, and mass arrays are created in the snapshot
    snapshot = pynbody.new(star=1,gas=nParticles)
    snapshot['vel'].units = v_unit
    snapshot['eps'] = 0.01*SimArray(np.ones(nParticles+1, dtype=np.float32), pos_unit)
    snapshot['metals'] = SimArray(np.zeros(nParticles+1, dtype=np.float32))
    snapshot['rho'] = SimArray(np.zeros(nParticles+1, dtype=np.float32))

    snapshot.gas['pos'] = xyz
    snapshot.gas['temp'] = ICobj.T(r)
    snapshot.gas['mass'] = m_particles
    snapshot.gas['metals'] = metals

    snapshot.star['pos'] = SimArray([[ 0.,  0.,  0.]],pos_unit)
    snapshot.star['vel'] = SimArray([[ 0.,  0.,  0.]], v_unit)
    snapshot.star['mass'] = m_star
    snapshot.star['metals'] = SimArray(star_metals)
    # Estimate the star's softening length as the closest particle distance
    #snapshot.star['eps'] = r.min()

    # Make param file
    param = isaac.make_param(snapshot, snapshotName)
    param['dMeanMolWeight'] = m

    gc.collect()

    # CALCULATE VELOCITY USING calc_velocity.py.  This also estimates the
    # gravitational softening length eps
    print 'Calculating circular velocity'
    preset = settings.changa_run.preset
    max_particles = global_settings['misc']['max_particles']
    calc_velocity.v_xy(snapshot, param, changa_preset=preset, max_particles=max_particles)

    gc.collect()

	# -------------------------------------------------
    # Estimate time step for changa to use
    # -------------------------------------------------
    # Save param file
    isaac.configsave(param, paramName, 'param')
    # Save snapshot
    snapshot.write(filename=snapshotName, fmt=pynbody.tipsy.TipsySnap)
    # est dDelta
    dDelta = ICgen_utils.est_time_step(paramName, preset)
    param['dDelta'] = dDelta

	# -------------------------------------------------
    # Create director file
    # -------------------------------------------------
    # largest radius to plot
    r_director = float(0.9 * r.max())
    # Maximum surface density
    sigma_min = float(ICobj.sigma(r_director))
    # surface density at largest radius
    sigma_max = float(ICobj.sigma.input_dict['sigma'].max())
    # Create director dict
    director = isaac.make_director(sigma_min, sigma_max, r_director, filename=param['achOutName'])
    ## Save .director file
    #isaac.configsave(director, directorName, 'director')

    #Now that velocities and everything are all initialized for gas particles, create new snapshot to return in which
    #single star particle is replaced by 2, same units as above
    snapshotBinary = pynbody.new(star=2,gas=nParticles)
    snapshotBinary['eps'] = 0.01*SimArray(np.ones(nParticles+2, dtype=np.float32), pos_unit)
    snapshotBinary['metals'] = SimArray(np.zeros(nParticles+2, dtype=np.float32))
    snapshotBinary['vel'].units = v_unit
    snapshotBinary['pos'].units = pos_unit
    snapshotBinary['mass'].units = snapshot['mass'].units
    snapshotBinary['rho'] = SimArray(np.zeros(nParticles+2, dtype=np.float32))

    #Assign gas particles with calculated/given values from above
    snapshotBinary.gas['pos'] = snapshot.gas['pos']
    snapshotBinary.gas['vel'] = snapshot.gas['vel']
    snapshotBinary.gas['temp'] = snapshot.gas['temp']
    snapshotBinary.gas['rho'] = snapshot.gas['rho']
    snapshotBinary.gas['eps'] = snapshot.gas['eps']
    snapshotBinary.gas['mass'] = snapshot.gas['mass']
    snapshotBinary.gas['metals'] = snapshot.gas['metals']

    #Load Binary system obj to initialize system
    binsys = ICobj.settings.physical.binsys

    x1,x2,v1,v2 = binsys.generateICs()

    #Put velocity in sim units
    #!!! Note: v_unit_vel will always be 29.785598165 km/s when m_unit = Msol and r_unit = 1 AU in kpc!!!
    #conv = v_unit_vel #km/s in sim units
    #v1 /= conv
    #v2 /= conv

    #Assign position, velocity assuming CCW orbit

    snapshotBinary.star[0]['pos'] = SimArray(x1,pos_unit)
    snapshotBinary.star[0]['vel'] = SimArray(v1,v_unit)
    snapshotBinary.star[1]['pos'] = SimArray(x2,pos_unit)
    snapshotBinary.star[1]['vel'] = SimArray(v2,v_unit)

    #Set stellar masses
    #Set Mass units
    #Create simArray for mass, convert units to simulation mass units
    priMass = SimArray(binsys.m1,m_unit)
    secMass = SimArray(binsys.m2,m_unit)

    snapshotBinary.star[0]['mass'] = priMass
    snapshotBinary.star[1]['mass'] = secMass
    snapshotBinary.star['metals'] = SimArray(star_metals)

    #Estimate stars' softening length as fraction of distance to COM
    #d = np.linalg.norm(x1) - np.linalg.norm(x2)
    d = np.sqrt(AddBinary.dotProduct(x1-x2,x1-x2))

    snapshotBinary.star[0]['eps'] = SimArray(math.fabs(d)/4.0,pos_unit)
    snapshotBinary.star[1]['eps'] = SimArray(math.fabs(d)/4.0,pos_unit)

    print 'Wrapping up'
    # Now set the star particle's tform to a negative number.  This allows
    # UW ChaNGa treat it as a sink particle.
    snapshotBinary.star['tform'] = -1.0

    #Set Sink Radius to be mass-weighted average of Roche lobes of two stars
    r1 = AddBinary.calcRocheLobe(binsys.m1/binsys.m2,binsys.a)
    r2 = AddBinary.calcRocheLobe(binsys.m2/binsys.m1,binsys.a)
    p = isaac.strip_units(binsys.m1/(binsys.m1 + binsys.m2))

    r_sink = (r1*p) + (r2*(1.0-p))
    param['dSinkBoundOrbitRadius'] = r_sink
    param['dSinkRadius'] = r_sink
    param['dSinkMassMin'] = 0.9 * isaac.strip_units(secMass)
    param['bDoSinks'] = 1

    return snapshotBinary, param, director

def make_director(ICobj, res=1200):

    director = {}
    director['render'] = 'tsc'
    director['FOV'] = 45.0
    director['clip'] = [0.0001, 500]
    director['up'] = [1, 0, 0]
    director['project'] = 'ortho'
    director['softgassph'] = 'softgassph'
    director['physical'] = 'physical'
    director['size'] = [res, res]

    sig_set = ICobj.settings.sigma
    mScale = ICobj.settings.snapshot.mScale
    snapshot_name = ICobj.settings.filenames.snapshotName
    f_prefix = os.path.splitext(os.path.basename(snapshot_name))[0]

    director['file'] = f_prefix

    if sig_set.kind == 'MQWS':

        rmax = sig_set.rout + 3*sig_set.rin
        zmax = float(rmax)
        director['eye'] = [0, 0, zmax]
        vmin = float(ICobj.rho(0, rmax))
        vmax = float(ICobj.rho.rho_binned[0,:].max())
        vmax *= mScale
        director['logscale'] = [vmin, 10*vmax]
        director['colgas'] = [1, 1, 1]

    return director

<commit_msg>final push before big merger<commit_after># -*- coding: utf-8 -*-
"""
Created on Fri Mar 21 15:11:31 2014

@author: ibackus
@editor: dfleming
-Note: indentation is 4 spaces in this file, not a tab!

"""

__version__ = "$Revision: 1 $"
# $Source$

import pynbody
SimArray = pynbody.array.SimArray
import numpy as np
import math
import gc
import os
import AddBinary
import isaac
import calc_velocity
import ICgen_utils
import ICglobal_settings
global_settings = ICglobal_settings.global_settings

def snapshot_gen(ICobj):
    """
    Generates a tipsy snapshot from the initial conditions object ICobj.

    Returns snapshot, param

        snapshot: tipsy snapshot
        param: dictionary containing info for a .param file
    Note: Code has been edited (dflemin3) such that now it returns a snapshot for a circumbinary disk
    where initial conditions generated assuming star at origin of mass M.  After gas initialized, replaced
    star at origin with binary system who's center of mass lies at the origin and who's mass m1 +m2 = M
    """

    print 'Generating snapshot...'
    # Constants
    G = SimArray(1.0,'G')
    # ------------------------------------
    # Load in things from ICobj
    # ------------------------------------
    print 'Accessing data from ICs'
    settings = ICobj.settings

    # snapshot file name
    snapshotName = settings.filenames.snapshotName
    paramName = settings.filenames.paramName

    # particle positions
    r = ICobj.pos.r
    xyz = ICobj.pos.xyz

    # Number of particles
    nParticles = ICobj.pos.nParticles

    # molecular mass
    m = settings.physical.m

    # star mass
    m_star = settings.physical.M.copy()

    # disk mass
    m_disk = ICobj.sigma.m_disk.copy()
    m_disk = isaac.match_units(m_disk, m_star)[0]

    # mass of the gas particles
    m_particles = m_disk / float(nParticles)

    # re-scale the particles (allows making of low-mass disk)
    m_particles *= settings.snapshot.mScale

    # -------------------------------------------------
    # Assign output
    # -------------------------------------------------
    print 'Assigning data to snapshot'
    # Get units all set up
    m_unit = m_star.units
    pos_unit = r.units

    if xyz.units != r.units:

        xyz.convert_units(pos_unit)

    # time units are sqrt(L^3/GM)
    t_unit = np.sqrt((pos_unit**3)*np.power((G*m_unit), -1)).units
    # velocity units are L/t
    v_unit = (pos_unit/t_unit).ratio('km s**-1')
    # Make it a unit, save value for future conversion
    v_unit_vel = v_unit
    #Ensure v_unit_vel is the same as what I assume it is.
    assert(np.fabs(AddBinary.VEL_UNIT-v_unit_vel)<AddBinary.SMALL),"VEL_UNIT not equal to ChaNGa unit! Why??"

    v_unit = pynbody.units.Unit('{0} km s**-1'.format(v_unit))

    # Other settings
    metals = settings.snapshot.metals
    star_metals = metals

    # Generate snapshot
    # Note that empty pos, vel, and mass arrays are created in the snapshot
    snapshot = pynbody.new(star=1,gas=nParticles)
    snapshot['vel'].units = v_unit
    snapshot['eps'] = 0.01*SimArray(np.ones(nParticles+1, dtype=np.float32), pos_unit)
    snapshot['metals'] = SimArray(np.zeros(nParticles+1, dtype=np.float32))
    snapshot['rho'] = SimArray(np.zeros(nParticles+1, dtype=np.float32))

    snapshot.gas['pos'] = xyz
    snapshot.gas['temp'] = ICobj.T(r)
    snapshot.gas['mass'] = m_particles
    snapshot.gas['metals'] = metals

    snapshot.star['pos'] = SimArray([[ 0.,  0.,  0.]],pos_unit)
    snapshot.star['vel'] = SimArray([[ 0.,  0.,  0.]], v_unit)
    snapshot.star['mass'] = m_star
    snapshot.star['metals'] = SimArray(star_metals)
    # Estimate the star's softening length as the closest particle distance
    #snapshot.star['eps'] = r.min()

    # Make param file
    param = isaac.make_param(snapshot, snapshotName)
    param['dMeanMolWeight'] = m

    gc.collect()

    # CALCULATE VELOCITY USING calc_velocity.py.  This also estimates the
    # gravitational softening length eps
    print 'Calculating circular velocity'
    preset = settings.changa_run.preset
    max_particles = global_settings['misc']['max_particles']
    calc_velocity.v_xy(snapshot, param, changa_preset=preset, max_particles=max_particles)

    gc.collect()

	# -------------------------------------------------
    # Estimate time step for changa to use
    # -------------------------------------------------
    # Save param file
    isaac.configsave(param, paramName, 'param')
    # Save snapshot
    snapshot.write(filename=snapshotName, fmt=pynbody.tipsy.TipsySnap)
    # est dDelta
    dDelta = ICgen_utils.est_time_step(paramName, preset)
    param['dDelta'] = dDelta

	# -------------------------------------------------
    # Create director file
    # -------------------------------------------------
    # largest radius to plot
    r_director = float(0.9 * r.max())
    # Maximum surface density
    sigma_min = float(ICobj.sigma(r_director))
    # surface density at largest radius
    sigma_max = float(ICobj.sigma.input_dict['sigma'].max())
    # Create director dict
    director = isaac.make_director(sigma_min, sigma_max, r_director, filename=param['achOutName'])
    ## Save .director file
    #isaac.configsave(director, directorName, 'director')

    #Now that velocities and everything are all initialized for gas particles, create new snapshot to return in which
    #single star particle is replaced by 2, same units as above
    snapshotBinary = pynbody.new(star=2,gas=nParticles)
    snapshotBinary['eps'] = 0.01*SimArray(np.ones(nParticles+2, dtype=np.float32), pos_unit)
    snapshotBinary['metals'] = SimArray(np.zeros(nParticles+2, dtype=np.float32))
    snapshotBinary['vel'].units = v_unit
    snapshotBinary['pos'].units = pos_unit
    snapshotBinary['mass'].units = snapshot['mass'].units
    snapshotBinary['rho'] = SimArray(np.zeros(nParticles+2, dtype=np.float32))

    #Assign gas particles with calculated/given values from above
    snapshotBinary.gas['pos'] = snapshot.gas['pos']
    snapshotBinary.gas['vel'] = snapshot.gas['vel']
    snapshotBinary.gas['temp'] = snapshot.gas['temp']
    snapshotBinary.gas['rho'] = snapshot.gas['rho']
    snapshotBinary.gas['eps'] = snapshot.gas['eps']
    snapshotBinary.gas['mass'] = snapshot.gas['mass']
    snapshotBinary.gas['metals'] = snapshot.gas['metals']

    #Load Binary system obj to initialize system
    binsys = ICobj.settings.physical.binsys

    x1,x2,v1,v2 = binsys.generateICs()

    #Put velocity in sim units
    #!!! Note: v_unit_vel will always be 29.785598165 km/s when m_unit = Msol and r_unit = 1 AU in kpc!!!
    #conv = v_unit_vel #km/s in sim units
    #v1 /= conv
    #v2 /= conv

    #Assign position, velocity assuming CCW orbit

    snapshotBinary.star[0]['pos'] = SimArray(x1,pos_unit)
    snapshotBinary.star[0]['vel'] = SimArray(v1,v_unit)
    snapshotBinary.star[1]['pos'] = SimArray(x2,pos_unit)
    snapshotBinary.star[1]['vel'] = SimArray(v2,v_unit)

    #Set stellar masses
    #Set Mass units
    #Create simArray for mass, convert units to simulation mass units
    priMass = SimArray(binsys.m1,m_unit)
    secMass = SimArray(binsys.m2,m_unit)

    snapshotBinary.star[0]['mass'] = priMass
    snapshotBinary.star[1]['mass'] = secMass
    snapshotBinary.star['metals'] = SimArray(star_metals)

    #Estimate stars' softening length as fraction of distance to COM
    d = np.sqrt(AddBinary.dotProduct(x1-x2,x1-x2))

    snapshotBinary.star[0]['eps'] = SimArray(math.fabs(d)/4.0,pos_unit)
    snapshotBinary.star[1]['eps'] = SimArray(math.fabs(d)/4.0,pos_unit)

    print 'Wrapping up'
    # Now set the star particle's tform to a negative number.  This allows
    # UW ChaNGa treat it as a sink particle.
    snapshotBinary.star['tform'] = -1.0

    #Set Sink Radius to be mass-weighted average of Roche lobes of two stars
    r1 = AddBinary.calcRocheLobe(binsys.m1/binsys.m2,binsys.a)
    r2 = AddBinary.calcRocheLobe(binsys.m2/binsys.m1,binsys.a)
    p = isaac.strip_units(binsys.m1/(binsys.m1 + binsys.m2))

    r_sink = (r1*p) + (r2*(1.0-p))
    param['dSinkBoundOrbitRadius'] = r_sink
    param['dSinkRadius'] = r_sink
    param['dSinkMassMin'] = 0.9 * isaac.strip_units(secMass)
    param['bDoSinks'] = 1

    return snapshotBinary, param, director

def make_director(ICobj, res=1200):

    director = {}
    director['render'] = 'tsc'
    director['FOV'] = 45.0
    director['clip'] = [0.0001, 500]
    director['up'] = [1, 0, 0]
    director['project'] = 'ortho'
    director['softgassph'] = 'softgassph'
    director['physical'] = 'physical'
    director['size'] = [res, res]

    sig_set = ICobj.settings.sigma
    mScale = ICobj.settings.snapshot.mScale
    snapshot_name = ICobj.settings.filenames.snapshotName
    f_prefix = os.path.splitext(os.path.basename(snapshot_name))[0]

    director['file'] = f_prefix

    if sig_set.kind == 'MQWS':

        rmax = sig_set.rout + 3*sig_set.rin
        zmax = float(rmax)
        director['eye'] = [0, 0, zmax]
        vmin = float(ICobj.rho(0, rmax))
        vmax = float(ICobj.rho.rho_binned[0,:].max())
        vmax *= mScale
        director['logscale'] = [vmin, 10*vmax]
        director['colgas'] = [1, 1, 1]

    return director

<|endoftext|>


<commit_before># This file is part of the MapProxy project.
# Copyright (C) 2010 Omniscale <http://omniscale.de>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
WMS clients for maps and information.
"""
from __future__ import with_statement
from mapproxy.config import base_config
from mapproxy.layer import MapQuery, InfoQuery
from mapproxy.source import SourceError
from mapproxy.client.http import HTTPClient
from mapproxy.srs import make_lin_transf
from mapproxy.image import ImageSource
from mapproxy.image.transform import ImageTransformer

class WMSClient(object):
    def __init__(self, request_template, supported_srs=None, http_client=None,
        resampling=None, supported_formats=None, lock=None):
        self.request_template = request_template
        self.http_client = http_client or HTTPClient()
        self.supported_srs = set(supported_srs or [])
        self.supported_formats = supported_formats or []
        self.resampling = resampling or base_config().image.resampling_method
        self.lock = lock

    def get_map(self, query):
        format = self.request_template.params.format
        if not format:
            format = query.format
            if self.supported_formats and format not in self.supported_formats:
                format = self.supported_formats[0]
        if self.supported_srs and query.srs not in self.supported_srs:
            return self._get_transformed(query, format)
        resp = self._retrieve(query, format)
        return ImageSource(resp, size=query.size, format=format)

    def _get_transformed(self, query, format):
        dst_srs = query.srs
        src_srs = self._best_supported_srs(dst_srs)
        dst_bbox = query.bbox
        src_bbox = dst_srs.transform_bbox_to(src_srs, dst_bbox)

        src_width, src_height = src_bbox[2]-src_bbox[0], src_bbox[3]-src_bbox[1]
        ratio = src_width/src_height
        dst_size = query.size
        xres, yres = src_width/dst_size[0], src_height/dst_size[1]
        if xres < yres:
            src_size = dst_size[0], int(dst_size[0]/ratio + 0.5)
        else:
            src_size = int(dst_size[1]*ratio +0.5), dst_size[1]

        src_query = MapQuery(src_bbox, src_size, src_srs, format)
        resp = self._retrieve(src_query, format)

        img = ImageSource(resp, format, size=src_size)

        img = ImageTransformer(src_srs, dst_srs, self.resampling).transform(img, src_bbox,
            query.size, dst_bbox)

        img.format = format
        return img

    def _best_supported_srs(self, srs):
        latlong = srs.is_latlong

        for srs in self.supported_srs:
            if srs.is_latlong == latlong:
                return srs

        return iter(self.supported_srs).next()

    def _retrieve(self, query, format):
        url = self._query_url(query, format)
        if self.lock:
            with self.lock():
                resp = self.http_client.open(url)
        else:
            resp = self.http_client.open(url)
        self._check_resp(resp)
        return resp

    def _check_resp(self, resp):
        if not resp.headers['Content-type'].startswith('image/'):
            raise SourceError('no image returned from source WMS')

    def _query_url(self, query, format):
        req = self.request_template.copy()
        req.params.bbox = query.bbox
        req.params.size = query.size
        req.params.srs = query.srs.srs_code
        req.params.format = format

        return req.complete_url

class WMSInfoClient(object):
    def __init__(self, request_template, supported_srs=None, http_client=None):
        self.request_template = request_template
        self.http_client = http_client or HTTPClient()
        self.supported_srs = set(supported_srs or [])

    def get_info(self, query):
        if self.supported_srs and query.srs not in self.supported_srs:
            return self._get_transformed(query)
        resp = self._retrieve(query)
        return resp

    def _get_transformed(self, query):
        req_srs = query.srs
        req_bbox = query.bbox
        info_srs = self._best_supported_srs(req_srs)
        info_bbox = req_srs.transform_bbox_to(info_srs, req_bbox)

        req_coord = make_lin_transf((0, query.size[1], query.size[0], 0), req_bbox)(query.pos)

        info_coord = req_srs.transform_to(info_srs, req_coord)
        info_pos = make_lin_transf((info_bbox), (0, query.size[1], query.size[0], 0))(info_coord)

        info_query = InfoQuery(info_bbox, query.size, info_srs, info_pos, query.info_format)
        return self._retrieve(info_query)

    def _best_supported_srs(self, srs):
        return iter(self.supported_srs).next()

    def _retrieve(self, query):
        url = self._query_url(query)
        return self.http_client.open(url)

    def _query_url(self, query):
        req = self.request_template.copy()
        req.params.bbox = query.bbox
        req.params.size = query.size
        req.params.pos = query.pos
        # del req.params['info_format']
        req.params['query_layers'] = req.params['layers']
        if query.info_format:
            req.params['info_format'] = query.info_format
        if not req.params.format:
            req.params.format = query.format or 'image/png'
        req.params.srs = query.srs.srs_code

        return req.complete_url
<commit_msg>WMSClient now always regards supported_formats<commit_after># This file is part of the MapProxy project.
# Copyright (C) 2010 Omniscale <http://omniscale.de>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
WMS clients for maps and information.
"""
from __future__ import with_statement
from mapproxy.config import base_config
from mapproxy.layer import MapQuery, InfoQuery
from mapproxy.source import SourceError
from mapproxy.client.http import HTTPClient
from mapproxy.srs import make_lin_transf
from mapproxy.image import ImageSource
from mapproxy.image.transform import ImageTransformer

class WMSClient(object):
    def __init__(self, request_template, supported_srs=None, http_client=None,
        resampling=None, supported_formats=None, lock=None):
        self.request_template = request_template
        self.http_client = http_client or HTTPClient()
        self.supported_srs = set(supported_srs or [])
        self.supported_formats = supported_formats or []
        self.resampling = resampling or base_config().image.resampling_method
        self.lock = lock

    def get_map(self, query):
        format = self.request_template.params.format
        if not format:
            format = query.format
        if self.supported_formats and format not in self.supported_formats:
            format = self.supported_formats[0]
        if self.supported_srs and query.srs not in self.supported_srs:
            return self._get_transformed(query, format)
        resp = self._retrieve(query, format)
        return ImageSource(resp, size=query.size, format=format)

    def _get_transformed(self, query, format):
        dst_srs = query.srs
        src_srs = self._best_supported_srs(dst_srs)
        dst_bbox = query.bbox
        src_bbox = dst_srs.transform_bbox_to(src_srs, dst_bbox)

        src_width, src_height = src_bbox[2]-src_bbox[0], src_bbox[3]-src_bbox[1]
        ratio = src_width/src_height
        dst_size = query.size
        xres, yres = src_width/dst_size[0], src_height/dst_size[1]
        if xres < yres:
            src_size = dst_size[0], int(dst_size[0]/ratio + 0.5)
        else:
            src_size = int(dst_size[1]*ratio +0.5), dst_size[1]

        src_query = MapQuery(src_bbox, src_size, src_srs, format)
        resp = self._retrieve(src_query, format)

        img = ImageSource(resp, format, size=src_size)

        img = ImageTransformer(src_srs, dst_srs, self.resampling).transform(img, src_bbox,
            query.size, dst_bbox)

        img.format = format
        return img

    def _best_supported_srs(self, srs):
        latlong = srs.is_latlong

        for srs in self.supported_srs:
            if srs.is_latlong == latlong:
                return srs

        return iter(self.supported_srs).next()

    def _retrieve(self, query, format):
        url = self._query_url(query, format)
        if self.lock:
            with self.lock():
                resp = self.http_client.open(url)
        else:
            resp = self.http_client.open(url)
        self._check_resp(resp)
        return resp

    def _check_resp(self, resp):
        if not resp.headers['Content-type'].startswith('image/'):
            raise SourceError('no image returned from source WMS')

    def _query_url(self, query, format):
        req = self.request_template.copy()
        req.params.bbox = query.bbox
        req.params.size = query.size
        req.params.srs = query.srs.srs_code
        req.params.format = format

        return req.complete_url

class WMSInfoClient(object):
    def __init__(self, request_template, supported_srs=None, http_client=None):
        self.request_template = request_template
        self.http_client = http_client or HTTPClient()
        self.supported_srs = set(supported_srs or [])

    def get_info(self, query):
        if self.supported_srs and query.srs not in self.supported_srs:
            return self._get_transformed(query)
        resp = self._retrieve(query)
        return resp

    def _get_transformed(self, query):
        req_srs = query.srs
        req_bbox = query.bbox
        info_srs = self._best_supported_srs(req_srs)
        info_bbox = req_srs.transform_bbox_to(info_srs, req_bbox)

        req_coord = make_lin_transf((0, query.size[1], query.size[0], 0), req_bbox)(query.pos)

        info_coord = req_srs.transform_to(info_srs, req_coord)
        info_pos = make_lin_transf((info_bbox), (0, query.size[1], query.size[0], 0))(info_coord)

        info_query = InfoQuery(info_bbox, query.size, info_srs, info_pos, query.info_format)
        return self._retrieve(info_query)

    def _best_supported_srs(self, srs):
        return iter(self.supported_srs).next()

    def _retrieve(self, query):
        url = self._query_url(query)
        return self.http_client.open(url)

    def _query_url(self, query):
        req = self.request_template.copy()
        req.params.bbox = query.bbox
        req.params.size = query.size
        req.params.pos = query.pos
        # del req.params['info_format']
        req.params['query_layers'] = req.params['layers']
        if query.info_format:
            req.params['info_format'] = query.info_format
        if not req.params.format:
            req.params.format = query.format or 'image/png'
        req.params.srs = query.srs.srs_code

        return req.complete_url
<|endoftext|>


<commit_before>class MarathonError(Exception):
    pass

class MarathonHttpError(MarathonError):

    def __init__(self, response):
        """
        :param :class:`requests.Response` response: HTTP response
        """
        content = response.json()
        self.status_code = response.status_code
        self.error_message = content.get('message')
        super(MarathonHttpError, self).__init__(self.__str__())

    def __repr__(self):
        return 'MarathonHttpError: HTTP %s returned with message, "%s"' % \
               (self.status_code, self.error_message)

    def __str__(self):
        return self.__repr__()

class NotFoundError(MarathonHttpError):
    pass

class InternalServerError(MarathonHttpError):
    pass

class InvalidChoiceError(MarathonError):

    def __init__(self, param, value, options):
        super(InvalidChoiceError, self).__init__(
            'Invalid choice "{value}" for param "{param}". Must be one of {options}'.format(
                param=param, value=value, options=options
            )
        )
<commit_msg>Handle HTTP errors without content graceful<commit_after>class MarathonError(Exception):
    pass

class MarathonHttpError(MarathonError):

    def __init__(self, response):
        """
        :param :class:`requests.Response` response: HTTP response
        """
        self.error_message = response.reason or ''
        if response.content:
            content = response.json()
            self.error_message = content.get('message', self.error_message)
        self.status_code = response.status_code
        super(MarathonHttpError, self).__init__(self.__str__())

    def __repr__(self):
        return 'MarathonHttpError: HTTP %s returned with message, "%s"' % \
               (self.status_code, self.error_message)

    def __str__(self):
        return self.__repr__()

class NotFoundError(MarathonHttpError):
    pass

class InternalServerError(MarathonHttpError):
    pass

class InvalidChoiceError(MarathonError):

    def __init__(self, param, value, options):
        super(InvalidChoiceError, self).__init__(
            'Invalid choice "{value}" for param "{param}". Must be one of {options}'.format(
                param=param, value=value, options=options
            )
        )
<|endoftext|>


<commit_before>import json

from django.conf import settings
from django.http import Http404, HttpResponse
from django.views.generic.base import View

from marimo.utils import smart_import

try:
    _marimo_widgets = smart_import(settings.MARIMO_REGISTRY)
except AttributeError:
    _marimo_widgets = {}

class MarimoRouter(View):
    """
    MarimoRouter splits up the request into individual widget packages and
    sends them to the handlers defined int he MARIMO_REGISTY
    """

    # TODO store widget registry as class variable so it can be overidden in instances

    def get(self, request):
        """ for a get request the bulk data is in request.GET """
        try:
            bulk = json.loads(request.GET['bulk'])
        except KeyError:
            raise Http404()
        else:
            return self.route(request, bulk)

    def route(self, request, bulk):
        """ this actually does the routing """
        response = []
        # TODO sanitize bulk
        for widget in bulk:
            print "doing: %s"%widget['id']
            # Try to get a callable from the dict... if it's not imported deal with it
            data = { 'id': widget['id'], }
            # TODO these try excepts are fucking hateful. Clean this up.
            try:
                view = _marimo_widgets[widget['widget_name']]
            except KeyError:
                print widget['widget_name']
                print _marimo_widgets
                data['status'] = 'WidgetNotFound'
            else:
                if not callable(view):
                    view = smart_import(view)()
                    _marimo_widgets[widget['widget_name']] = view

                try:
                    data.update(view(request, *widget['args'], **widget['kwargs']))
                except Exception, e:
                    try:
                        data = view.on_error(e, data, request, *widget['args'], **widget['kwargs'])
                    except:
                        # BARE EXCEPPPPTTTT TODO TODO TODO
                        data = view.on_error(e, data, request, *[], **{})
                else:
                    data['status'] = 'succeeded'
            finally:
                response.append(data)

        return self.build_response(request, response)

    def build_response(self, request, data):
        as_json = json.dumps(data)
        if request.REQUEST.get('format') == 'jsonp' and request.REQUEST.get('callback'):
            mimetype = 'text/javascript'
            callback = request.REQUEST.get('callback')
            as_json = "{0}({1});".format(callback, as_json)
        else:
            mimetype = 'application/json'

        return HttpResponse(as_json, mimetype=mimetype)
<commit_msg>Revert "temporary hack, revert me"<commit_after>import json

from django.conf import settings
from django.http import Http404, HttpResponse
from django.views.generic.base import View

from marimo.utils import smart_import

try:
    _marimo_widgets = smart_import(settings.MARIMO_REGISTRY)
except AttributeError:
    _marimo_widgets = {}

class MarimoRouter(View):
    """
    MarimoRouter splits up the request into individual widget packages and
    sends them to the handlers defined int he MARIMO_REGISTY
    """

    # TODO store widget registry as class variable so it can be overidden in instances

    def get(self, request):
        """ for a get request the bulk data is in request.GET """
        try:
            bulk = json.loads(request.GET['bulk'])
        except KeyError:
            raise Http404()
        else:
            return self.route(request, bulk)

    def route(self, request, bulk):
        """ this actually does the routing """
        response = []
        # TODO sanitize bulk
        for widget in bulk:
            print "doing: %s"%widget['id']
            # Try to get a callable from the dict... if it's not imported deal with it
            data = { 'id': widget['id'], }
            try:
                view = _marimo_widgets[widget['widget_name']]
            except KeyError:
                print widget['widget_name']
                print _marimo_widgets
                data['status'] = 'WidgetNotFound'
            else:
                if not callable(view):
                    view = smart_import(view)()
                    _marimo_widgets[widget['widget_name']] = view

                try:
                    data.update(view(request, *widget['args'], **widget['kwargs']))
                except Exception, e:
                    data = view.on_error(e, data, request, *widget['args'], **widget['kwargs'])
                else:
                    data['status'] = 'succeeded'
            finally:
                response.append(data)

        return self.build_response(request, response)

    def build_response(self, request, data):
        as_json = json.dumps(data)
        if request.REQUEST.get('format') == 'jsonp' and request.REQUEST.get('callback'):
            mimetype = 'text/javascript'
            callback = request.REQUEST.get('callback')
            as_json = "{0}({1});".format(callback, as_json)
        else:
            mimetype = 'application/json'

        return HttpResponse(as_json, mimetype=mimetype)
<|endoftext|>


<commit_before># encoding: utf-8

from copy import copy

from marrow.util.object import NoDefault
from marrow.tags import html5 as tag

from transforms import BaseTransform, BooleanTransform

__all__ = ['Widget', 'NestedWidget', 'Form', 'FieldSet', 'Label', 'Layout', 'Input', 'BooleanInput', 'Link']

class Widget(object):
    transform = BaseTransform()
    default = None

    def __init__(self, name_, title_=None, transform=NoDefault, default=NoDefault, data_=NoDefault, **kw):
        self.name = name_
        self.title = title_
        self.data = dict() if data_ is NoDefault else data_
        self.args = kw

        if transform is not NoDefault: self.transform = transform
        if default is not NoDefault: self.default = default

    @property
    def value(self):
        value = self.data.get(self.name, self.default)
        return self.transform(value) if self.transform else value

    def native(self, data):
        value = data.get(self.name, None)

        if value is None:
            return self.default

        return self.transform.native(value) if self.transform else value

    @property
    def template(self):
        raise NotImplementedError

    def __call__(self, data=NoDefault):
        local = copy(self) # Thread Safety

        if data is not NoDefault:
            if isinstance(local.data, dict):
                local.data.update(data)

            else:
                local.data = data

        return local.template

class NestedWidget(Widget):
    def __init__(self, name_, title_=None, children=None, *args, **kw):
        self.children = children if children else list()
        super(NestedWidget, self).__init__(name_, title_, *args, **kw)

    def native(self, data):
        result = dict()

        for child in self.children:
            if isinstance(child, NestedWidget):
                result.update(child.native(data)[0])
                continue

            result[child.name] = child.native(data)

        remaining = dict()

        for i in data:
            if i not in result:
                remaining[i] = data[i]

        return result, remaining

class Form(NestedWidget):
    def __init__(self, name_, title_=None, layout=None, children=[], footer=None, *args, **kw):
        super(Form, self).__init__(name_, title_, children, *args, **kw)

        self.layout = layout
        self.footer = footer

        if 'method' not in self.args:
            self.args['method'] = 'post'

    @property
    def template(self):
        return tag.form ( id = self.name + '-form', **self.args ) [
                ([
                    self.layout(self.name, children=self.children)(self.data)
                ] if self.layout else [
                    child(self.data) for child in self.children
                ]) + ([
                    self.footer(self.name)(self.data) if isinstance(self.footer, type) else self.footer(self.data)
                ] if self.footer else [])
            ]

class FieldSet(NestedWidget):
    def __init__(self, name_, title_=None, layout=None, children=[], *args, **kw):
        super(FieldSet, self).__init__(name_, title_, children, *args, **kw)

        self.layout = layout

    @property
    def template(self):
        return tag.fieldset (
                id = self.name + '-set',
                **self.args
            ) [
                ([tag.legend [ self.title ] ] if self.title else []) + ([
                    self.layout(self.name, children=self.children)(self.data)
                ] if self.layout else [
                    child(self.data) for child in self.children
                ])
            ]

class Label(Widget):
    def __init__(self, name_, title_=None, for_=None, *args, **kw):
        self.for_ = for_
        super(Label, self).__init__(name_, title_, *args, **kw)

    @property
    def template(self):
        return tag.label (
                for_ = (self.for_.name + '-field') if self.for_ else None
            ) [ self.title if self.title else self.for_.title ]

class Layout(NestedWidget):
    label = Label

class Input(Widget):
    type_ = None

    @property
    def template(self):
        return tag.input (
                type_ = self.type_,
                name = self.name,
                id = self.name + '-field',
                value = self.value,
                **self.args
            )

class BooleanInput(Input):
    transform = BooleanTransform()

    @property
    def template(self):
        return tag.div [[ tag.input (
                    type_ = self.type_,
                    name = self.name,
                    id = self.name + '-field',
                    checked = self.value,
                    **self.args
                )] + [
                    tag.label ( for_ = self.name + '-field' ) [ self.args.get('title') ]
                ] if 'title' in self.args else []]

class Link(Widget):
    @property
    def template(self):
        return tag.a ( id = self.name + "-link", **self.args ) [ self.title ]
<commit_msg>Update to allow use of MongoEngine objects (which support dictionary access), not just dicts.<commit_after># encoding: utf-8

from copy import copy

from marrow.util.object import NoDefault
from marrow.tags import html5 as tag

from transforms import BaseTransform, BooleanTransform

__all__ = ['Widget', 'NestedWidget', 'Form', 'FieldSet', 'Label', 'Layout', 'Input', 'BooleanInput', 'Link']

class Widget(object):
    transform = BaseTransform()
    default = None

    def __init__(self, name_, title_=None, transform=NoDefault, default=NoDefault, data_=NoDefault, **kw):
        self.name = name_
        self.title = title_
        self.data = dict() if data_ is NoDefault else data_
        self.args = kw

        if transform is not NoDefault: self.transform = transform
        if default is not NoDefault: self.default = default

    @property
    def value(self):
        value = self.data.get(self.name, self.default)
        return self.transform(value) if self.transform else value

    def native(self, data):
        value = data.get(self.name, None)

        if value is None:
            return self.default

        return self.transform.native(value) if self.transform else value

    @property
    def template(self):
        raise NotImplementedError

    def __call__(self, data=NoDefault):
        local = copy(self) # Thread Safety

        if data is not NoDefault:
            if isinstance(local.data, dict) and isinstance(data, dict):
                local.data.update(data)

            else:
                local.data = data

        return local.template

class NestedWidget(Widget):
    def __init__(self, name_, title_=None, children=None, *args, **kw):
        self.children = children if children else list()
        super(NestedWidget, self).__init__(name_, title_, *args, **kw)

    def native(self, data):
        result = dict()

        for child in self.children:
            if isinstance(child, NestedWidget):
                result.update(child.native(data)[0])
                continue

            result[child.name] = child.native(data)

        remaining = dict()

        for i in data:
            if i not in result:
                remaining[i] = data[i]

        return result, remaining

class Form(NestedWidget):
    def __init__(self, name_, title_=None, layout=None, children=[], footer=None, *args, **kw):
        super(Form, self).__init__(name_, title_, children, *args, **kw)

        self.layout = layout
        self.footer = footer

        if 'method' not in self.args:
            self.args['method'] = 'post'

    @property
    def template(self):
        return tag.form ( id = self.name + '-form', **self.args ) [
                ([
                    self.layout(self.name, children=self.children)(self.data)
                ] if self.layout else [
                    child(self.data) for child in self.children
                ]) + ([
                    self.footer(self.name)(self.data) if isinstance(self.footer, type) else self.footer(self.data)
                ] if self.footer else [])
            ]

class FieldSet(NestedWidget):
    def __init__(self, name_, title_=None, layout=None, children=[], *args, **kw):
        super(FieldSet, self).__init__(name_, title_, children, *args, **kw)

        self.layout = layout

    @property
    def template(self):
        return tag.fieldset (
                id = self.name + '-set',
                **self.args
            ) [
                ([tag.legend [ self.title ] ] if self.title else []) + ([
                    self.layout(self.name, children=self.children)(self.data)
                ] if self.layout else [
                    child(self.data) for child in self.children
                ])
            ]

class Label(Widget):
    def __init__(self, name_, title_=None, for_=None, *args, **kw):
        self.for_ = for_
        super(Label, self).__init__(name_, title_, *args, **kw)

    @property
    def template(self):
        return tag.label (
                for_ = (self.for_.name + '-field') if self.for_ else None
            ) [ self.title if self.title else self.for_.title ]

class Layout(NestedWidget):
    label = Label

class Input(Widget):
    type_ = None

    @property
    def template(self):
        return tag.input (
                type_ = self.type_,
                name = self.name,
                id = self.name + '-field',
                value = self.value,
                **self.args
            )

class BooleanInput(Input):
    transform = BooleanTransform()

    @property
    def template(self):
        return tag.div [[ tag.input (
                    type_ = self.type_,
                    name = self.name,
                    id = self.name + '-field',
                    checked = self.value,
                    **self.args
                )] + [
                    tag.label ( for_ = self.name + '-field' ) [ self.args.get('title') ]
                ] if 'title' in self.args else []]

class Link(Widget):
    @property
    def template(self):
        return tag.a ( id = self.name + "-link", **self.args ) [ self.title ]
<|endoftext|>


<commit_before>import discord
import re
import csv
import string
import random
from discord.ext import commands

class MCOCTools:
    '''Tools for Marvel Contest of Champions'''
    lookup_links = {
            'event': (
                'Tiny MCoC Schedule',
                '<http://simians.tk/MCOC-Sched>',
                'Josh Morris Schedule',
                'https://d2jixqqjqj5d23.cloudfront.net/assets/developer/imgs/icons/google-spreadsheet-icon.png'),
            'hook': (
                'hook/champions',
                '<http://hook.github.io/champions>',
                'hook/champions for Collector',
                'https://assets-cdn.github.com/favicon.ico'),
            'spotlight': (
                'MCoC Spotlight',
                '<http://simians.tk/MCoCspotlight>'),
            'marvelsynergy': (
                'Team Builder',
                '<http://www.marvelsynergy.com/team-builder>',
                'Marvel Synergy',
                'http://www.marvelsynergy.com/images/marvelsynergy.png'),
            'alsciende':(
                'Alsciende Mastery Tool',
                '<https://alsciende.github.io/masteries/v10.0.1/#>',
                'by u/alsciende',
                'https://assets-cdn.github.com/favicon.ico'),
            'simulator': (
                'Mastery Simulator',
                '<http://simians.tk/msimSDF>'),
            'streak': (
                'Infinite Streak',
                'http://simians.tk/-sdf-streak')
                #'http://simians.tk/SDFstreak')
    }
    mcolor = discord.Color.gold()
    icon_sdf = 'https://raw.githubusercontent.com/JasonJW/mcoc-cogs/master/mcoc/data/sdf_icon.png'
    dataset = 'data/mcoc/masteries.csv'

    def __init__(self, bot):
        self.bot = bot

    def present(self, lookup):
        em=discord.Embed(color=self.mcolor,title=lookup[0],description=lookup[1])
        print(len(lookup))
        if len(lookup) > 2:
            em.set_footer(text=lookup[2],icon_url=lookup[3])
        else:
            em.set_footer(text='Presented by [-SDF-]',icon_url=self.icon_sdf)
        return em

    #Gold Realm disabled
    # @commands.command()
    # async def gold(self,msg = None):
    #     '''Gold Realm Schedule'''
    #     package = '```           PST      CST      EST\nJun 5    10:00    12:00    13:00\nJun 5    19:00    21:00    22:00\nJun 6     4:00     6:00     7:00\nJun 6    13:00    15:00    16:00\nJun 6    22:00     0:00     1:00\nJun 7     7:00     9:00    10:00\nJun 7    16:00    18:00    19:00\nJun 8     1:00     3:00     4:00\nJun 8    10:00    12:00    13:00\nJun 8    19:00    21:00    22:00\nJun 9     4:00     6:00     7:00\nJun 9    13:00    15:00    16:00\nJun 10   22:00     0:00     1:00\nJun 11    7:00     9:00    10:00\nJun 11   16:00    18:00    19:00\nJun 12    1:00     3:00     4:00```'
    #     # await self.bot.say('```'+package+'```')
    #     if msg is not None:
    #         package = msg+'\n'+package
    #     em = discord.Embed(color=discord.Color.gold(),title='Gold Realm Schedule',description=package)
    #     em.set_footer(text='Presented by [-SDF-]',icon_url=self.icon_sdf)
    #     await self.bot.say(embed=em)

    @commands.command(help=lookup_links['event'][0], aliases=['events','schedule',])
    async def event(self):
        lookup = self.lookup_links['event']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['spotlight'][0],)
    async def spotlight(self):
        lookup = self.lookup_links['spotlight']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['marvelsynergy'][0])
    async def marvelsynergy(self):
        lookup = self.lookup_links['marvelsynergy']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['simulator'][0],aliases=['msim'])
    async def simulator(self):
        lookup = self.lookup_links['simulator']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['alsciende'][0], aliases=('mrig',))
    async def alsciende(self):
        lookup = self.lookup_links['alsciende']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['streak'][0])
    async def streak(self):
        lookup = self.lookup_links['streak']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['hook'][0])
    async def hook(self):
        lookup = self.lookup_links['hook']
        await self.bot.say(embed=self.present(lookup))

    @commands.command()
    asynce def keygen(self, prefix='SECEMP'):
        string.letters={'ABCDEFGHIJKLMNOPQURSTUVWXYZ'}
        string.numbers={'0123456789'}
        package = []
        for i in range{0,10}:
            nums='{}{}'.format(random.choice(string.numbers),random.choice(string.numbers))
            letters='{}{}{}{}'.format(random.choic(string.letters),random.choic(string.letters),random.choic(string.letters),random.choic(string.letters))
            package.append(prefix+nums+letters)
        em=discord.Embed(color=discord.Color.gold(),title='Email Code Generator',description='\n'.join(package))
        await self.bot.say(embed=em)
#### Mastery stuff###
# class MCOCMastery:
    masteryColor={'Offense': discord.Color.red(),
            'Defense': discord.Color.blue(),
            'Utility': discord.Color.green()}
    tokens={'Vitality','Greater Vitality','Salve','Recovery','Energy Resistance','Physical Resistance','Block Proficiency','Perfect Block','Stand Your Ground','Collar Tech','Serum Science','Willpower','Coagulate','Suture','Inequity','Resonate',
            'Strength','Greater Strength','Pierce','Lesser Precision','Precision','Lesser Cruelty','Cruelty','Courage','Extended Fury','Enhanced Fury','Pure Skill','Mutagenesis','Glass Canon','Recoil','Liquid Courage','Double Edge','Despair','DeepWounds','Unfazed','Assassin',
            'Wisdom','Intelligence','Limber','Parry','Stupefy','Petrify','Pacify','Dexterity','Pitance','Prosperity','Cosmic Awareness','Mystic Dispersion','Detect Cosmic','Detect Tech','Detect Mystic','Scouter Lens','Detect Mutant','Detect Science','Detect Skill'}

    @commands.group(pass_context=True, aliases=['masteries',])
    async def mastery(self, ctx):
        if ctx.invoked_subcommand is None:
            await self.bot.send_cmd_help(ctx)
            return

    @mastery.command(pass_context=True,)
    async def test(self,ctx):
        message = ctx.message.content
        # mastery='Resonate'
        # rank=3
        # rows = get_csv_rows(self.dataset,'Mastery',mastery)
        #  text = []
        # cost = {'ucarb': 'Carbonadium Mastery Core', 'ustony': 'Stony Mastery Core', 'uclass': '', 'uunit': 'Units', 'rgold': 'Gold', 'runit': 'Units'}
        # cores = {'Collar Tech': 'Tech Core', 'Serum Science': 'Mastery Serum', 'Mutagenesis': 'Mastery Core X', 'Pure Skill': 'Mastery Core of Apptitude', 'Cosmic Awareness':'Cosmic Mastery Core', 'Mystic Dispersion': 'Mystical Mastery Core',
        #         'Detect Tech': 'Tech Core', 'Detect Science': 'Mastery Serum', 'Detect Mutant': 'Mastery Core X', 'Detect Skill': 'Mastery Core of Apptitude', 'Detect Cosmic':'Cosmic Mastery Core', 'Detect Mystic': 'Mystical Mastery Core',}
        # if mastery in cores:
        #     cost['uclass'][1]=cores[mastery]
        for token in self.tokens:
            print(token+' found')
            if token in message:
                maxrank = rows[0]['Max Ranks']
                category = rows[0]['Category']

                em=discord.Embed(color=self.masteryColor[category],title=mastery,description=' '.join(t for t in text))
                # for r in maxrank:
                    # unlock, rankcost = _get_cost(mastery, r)
                unlock, rankcost = self._get_cost(mastery, maxrank)

                em.add_field(name='Unlock Cost',value='\n'.join(u for u in unlock))
                em.add_field(name='Rank Cost',value='\n'.join(r for r in rankcost))
                await self.bot.say(embed=em)

    @mastery.command(pass_context=True, name='cost')
    async def _cost(self,ctx):
        message = ctx.message.content.split(' ')
        stack = []
        for i in range(0,len(message)-1):
            if i+1 < len(message):
                test='{} {}'.format(message[i],message[i+1])
                if test in self.tokens:
                    stack.append(test)
                    i=i+2
                elif message[i] in self.tokens:
                    stack.append(message[i])
            else:
                if message[i] in self.tokens:
                    stack.append(message[i])
        await self.bot.say('\n'.join(stack))

        # args = ctx.split(' ')
        #
        # for arg in args:
        #     # parse_re = re.compile(r'''(?:r(?P<rank>[0-9]))|(?:d(?P<debug>[0-9]{1,2}))''', re.X)
        #     if arg in tokens:
        #         row = mcoc.csv_get_row(mcoc.data_files['masteries']['local'], 'MasteryExt', arg)

        # await self.bot.say('Dummy message for cost')

    @mastery.command(pass_context=True, name='set')
    async def _set(self,ctx):
        await self.bot.say('Dummy message for set')

    def _get_cost(self, mastery, rank):
        row = csv_get_row(self.dataset,'Mastery',mastery)
        cost = {'ucarb': 'Carbonadium Mastery Core', 'ustony': 'Stony Mastery Core', 'uclass': '', 'uunit': 'Units', 'rgold': 'Gold', 'runit': 'Units'}
        cores = {'Collar Tech': 'Tech Core', 'Serum Science': 'Mastery Serum', 'Mutagenesis': 'Mastery Core X', 'Pure Skill': 'Mastery Core of Apptitude', 'Cosmic Awareness':'Cosmic Mastery Core', 'Mystic Dispersion': 'Mystical Mastery Core',
                'Detect Tech': 'Tech Core', 'Detect Science': 'Mastery Serum', 'Detect Mutant': 'Mastery Core X', 'Detect Skill': 'Mastery Core of Apptitude', 'Detect Cosmic':'Cosmic Mastery Core', 'Detect Mystic': 'Mystical Mastery Core',}
        if mastery in cores:
            cost['uclass'][1]=cores[mastery]
        unlock=[]
        rankcost=[]
        for c in {'ucarb','ustony','uclass','uunit',}:
            key = '{}{}'.format(c, rank)
            print(key)
            price = rows[0][key]
            print(price)
            if price != '':
                if int(price) > 0:
                    core = cost[c]
                    unlock.append('{0}x {1}'.format(price, core))
        for c in {'rgold','runit',}:
            key = '{}{}'.format(c, rank)
            price = rows[0][key]
            print(price)
            if price != '':
                if int(price) > 0:
                    core = cost[c]
                    rankcost.append('{0}x {1}'.format(price, core))

        return unlock, rankcost

    def _get_text(self, mastery, rank):
        rows = csv_get_rows(self.dataset,'Mastery',mastery)
        for row in rows:
            text.append(row['Text'].format(row[str(rank)]))
        return text

def load_csv(filename):
    return csv.DictReader(open(filename))

def get_csv_row(filecsv, column, match_val, default=None):
    print(match_val)
    csvfile = load_csv(filecsv)
    for row in csvfile:
        if row[column] == match_val:
            if default is not None:
                for k, v in row.items():
                    if v == '':
                        row[k] = default
            return row

def get_csv_rows(filecsv, column, match_val, default=None):
    print(match_val)
    csvfile = load_csv(filecsv)
    package =[]
    for row in csvfile:
        if row[column] == match_val:
            if default is not None:
                for k, v in row.items():
                    if v == '':
                        row[k] = default
            package.append(row)
    return package

def setup(bot):
    bot.add_cog(MCOCTools(bot))
<commit_msg>re-think mastery cost<commit_after>import discord
import re
import csv
import string
import random
from discord.ext import commands

class MCOCTools:
    '''Tools for Marvel Contest of Champions'''
    lookup_links = {
            'event': (
                'Tiny MCoC Schedule',
                '<http://simians.tk/MCOC-Sched>',
                'Josh Morris Schedule',
                'https://d2jixqqjqj5d23.cloudfront.net/assets/developer/imgs/icons/google-spreadsheet-icon.png'),
            'hook': (
                'hook/champions',
                '<http://hook.github.io/champions>',
                'hook/champions for Collector',
                'https://assets-cdn.github.com/favicon.ico'),
            'spotlight': (
                'MCoC Spotlight',
                '<http://simians.tk/MCoCspotlight>'),
            'marvelsynergy': (
                'Team Builder',
                '<http://www.marvelsynergy.com/team-builder>',
                'Marvel Synergy',
                'http://www.marvelsynergy.com/images/marvelsynergy.png'),
            'alsciende':(
                'Alsciende Mastery Tool',
                '<https://alsciende.github.io/masteries/v10.0.1/#>',
                'by u/alsciende',
                'https://assets-cdn.github.com/favicon.ico'),
            'simulator': (
                'Mastery Simulator',
                '<http://simians.tk/msimSDF>'),
            'streak': (
                'Infinite Streak',
                'http://simians.tk/-sdf-streak')
                #'http://simians.tk/SDFstreak')
    }
    mcolor = discord.Color.gold()
    icon_sdf = 'https://raw.githubusercontent.com/JasonJW/mcoc-cogs/master/mcoc/data/sdf_icon.png'
    dataset = 'data/mcoc/masteries.csv'

    def __init__(self, bot):
        self.bot = bot

    def present(self, lookup):
        em=discord.Embed(color=self.mcolor,title=lookup[0],description=lookup[1])
        print(len(lookup))
        if len(lookup) > 2:
            em.set_footer(text=lookup[2],icon_url=lookup[3])
        else:
            em.set_footer(text='Presented by [-SDF-]',icon_url=self.icon_sdf)
        return em

    #Gold Realm disabled
    # @commands.command()
    # async def gold(self,msg = None):
    #     '''Gold Realm Schedule'''
    #     package = '```           PST      CST      EST\nJun 5    10:00    12:00    13:00\nJun 5    19:00    21:00    22:00\nJun 6     4:00     6:00     7:00\nJun 6    13:00    15:00    16:00\nJun 6    22:00     0:00     1:00\nJun 7     7:00     9:00    10:00\nJun 7    16:00    18:00    19:00\nJun 8     1:00     3:00     4:00\nJun 8    10:00    12:00    13:00\nJun 8    19:00    21:00    22:00\nJun 9     4:00     6:00     7:00\nJun 9    13:00    15:00    16:00\nJun 10   22:00     0:00     1:00\nJun 11    7:00     9:00    10:00\nJun 11   16:00    18:00    19:00\nJun 12    1:00     3:00     4:00```'
    #     # await self.bot.say('```'+package+'```')
    #     if msg is not None:
    #         package = msg+'\n'+package
    #     em = discord.Embed(color=discord.Color.gold(),title='Gold Realm Schedule',description=package)
    #     em.set_footer(text='Presented by [-SDF-]',icon_url=self.icon_sdf)
    #     await self.bot.say(embed=em)

    @commands.command(help=lookup_links['event'][0], aliases=['events','schedule',])
    async def event(self):
        lookup = self.lookup_links['event']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['spotlight'][0],)
    async def spotlight(self):
        lookup = self.lookup_links['spotlight']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['marvelsynergy'][0])
    async def marvelsynergy(self):
        lookup = self.lookup_links['marvelsynergy']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['simulator'][0],aliases=['msim'])
    async def simulator(self):
        lookup = self.lookup_links['simulator']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['alsciende'][0], aliases=('mrig',))
    async def alsciende(self):
        lookup = self.lookup_links['alsciende']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['streak'][0])
    async def streak(self):
        lookup = self.lookup_links['streak']
        await self.bot.say(embed=self.present(lookup))

    @commands.command(help=lookup_links['hook'][0])
    async def hook(self):
        lookup = self.lookup_links['hook']
        await self.bot.say(embed=self.present(lookup))

    @commands.command()
    asynce def keygen(self, prefix='SECEMP'):
        string.letters={'ABCDEFGHIJKLMNOPQURSTUVWXYZ'}
        string.numbers={'0123456789'}
        package = []
        for i in range{0,10}:
            nums='{}{}'.format(random.choice(string.numbers),random.choice(string.numbers))
            letters='{}{}{}{}'.format(random.choic(string.letters),random.choic(string.letters),random.choic(string.letters),random.choic(string.letters))
            package.append(prefix+nums+letters)
        em=discord.Embed(color=discord.Color.gold(),title='Email Code Generator',description='\n'.join(package))
        await self.bot.say(embed=em)
#### Mastery stuff###
# class MCOCMastery:
    masteryColor={'Offense': discord.Color.red(),
            'Defense': discord.Color.blue(),
            'Utility': discord.Color.green()}
    tokens={'Vitality','Greater Vitality','Salve','Recovery','Energy Resistance','Physical Resistance','Block Proficiency','Perfect Block','Stand Your Ground','Collar Tech','Serum Science','Willpower','Coagulate','Suture','Inequity','Resonate',
            'Strength','Greater Strength','Pierce','Lesser Precision','Precision','Lesser Cruelty','Cruelty','Courage','Extended Fury','Enhanced Fury','Pure Skill','Mutagenesis','Glass Canon','Recoil','Liquid Courage','Double Edge','Despair','DeepWounds','Unfazed','Assassin',
            'Wisdom','Intelligence','Limber','Parry','Stupefy','Petrify','Pacify','Dexterity','Pitance','Prosperity','Cosmic Awareness','Mystic Dispersion','Detect Cosmic','Detect Tech','Detect Mystic','Scouter Lens','Detect Mutant','Detect Science','Detect Skill'}

    @commands.group(pass_context=True, aliases=['masteries',])
    async def mastery(self, ctx):
        if ctx.invoked_subcommand is None:
            await self.bot.send_cmd_help(ctx)
            return

    @mastery.command(pass_context=True,)
    async def test(self,ctx):
        message = ctx.message.content
        # mastery='Resonate'
        # rank=3
        # rows = get_csv_rows(self.dataset,'Mastery',mastery)
        #  text = []
        # cost = {'ucarb': 'Carbonadium Mastery Core', 'ustony': 'Stony Mastery Core', 'uclass': '', 'uunit': 'Units', 'rgold': 'Gold', 'runit': 'Units'}
        # cores = {'Collar Tech': 'Tech Core', 'Serum Science': 'Mastery Serum', 'Mutagenesis': 'Mastery Core X', 'Pure Skill': 'Mastery Core of Apptitude', 'Cosmic Awareness':'Cosmic Mastery Core', 'Mystic Dispersion': 'Mystical Mastery Core',
        #         'Detect Tech': 'Tech Core', 'Detect Science': 'Mastery Serum', 'Detect Mutant': 'Mastery Core X', 'Detect Skill': 'Mastery Core of Apptitude', 'Detect Cosmic':'Cosmic Mastery Core', 'Detect Mystic': 'Mystical Mastery Core',}
        # if mastery in cores:
        #     cost['uclass'][1]=cores[mastery]
        for token in self.tokens:
            print(token+' found')
            if token in message:
                maxrank = rows[0]['Max Ranks']
                category = rows[0]['Category']

                em=discord.Embed(color=self.masteryColor[category],title=mastery,description=' '.join(t for t in text))
                # for r in maxrank:
                    # unlock, rankcost = _get_cost(mastery, r)
                unlock, rankcost = self._get_cost(mastery, maxrank)

                em.add_field(name='Unlock Cost',value='\n'.join(u for u in unlock))
                em.add_field(name='Rank Cost',value='\n'.join(r for r in rankcost))
                await self.bot.say(embed=em)

    @mastery.command(pass_context=True, name='cost')
    async def _cost(self,ctx):
        # message = ctx.message.content.split(' ')
        # stack = []
        # for i in range(0,len(message)-1):
        #     if i+1 < len(message):
        #         test='{} {}'.format(message[i],message[i+1])
        #         if test in self.tokens:
        #             stack.append(test)
        #             i=i+2
        #         elif message[i] in self.tokens:
        #             stack.append(message[i])
        #     else:
        #         if message[i] in self.tokens:
        #             stack.append(message[i])
        # await self.bot.say('\n'.join(stack))

        # args = ctx.split(' ')
        #
        # for arg in args:
        #     # parse_re = re.compile(r'''(?:r(?P<rank>[0-9]))|(?:d(?P<debug>[0-9]{1,2}))''', re.X)
        #     if arg in tokens:
        #         row = mcoc.csv_get_row(mcoc.data_files['masteries']['local'], 'MasteryExt', arg)

        await self.bot.say('Dummy message for cost')

    @mastery.command(pass_context=True, name='set')
    async def _set(self,ctx):
        await self.bot.say('Dummy message for set')

    def _get_cost(self, mastery, rank):
        row = csv_get_row(self.dataset,'Mastery',mastery)
        cost = {'ucarb': 'Carbonadium Mastery Core', 'ustony': 'Stony Mastery Core', 'uclass': '', 'uunit': 'Units', 'rgold': 'Gold', 'runit': 'Units'}
        cores = {'Collar Tech': 'Tech Core', 'Serum Science': 'Mastery Serum', 'Mutagenesis': 'Mastery Core X', 'Pure Skill': 'Mastery Core of Apptitude', 'Cosmic Awareness':'Cosmic Mastery Core', 'Mystic Dispersion': 'Mystical Mastery Core',
                'Detect Tech': 'Tech Core', 'Detect Science': 'Mastery Serum', 'Detect Mutant': 'Mastery Core X', 'Detect Skill': 'Mastery Core of Apptitude', 'Detect Cosmic':'Cosmic Mastery Core', 'Detect Mystic': 'Mystical Mastery Core',}
        if mastery in cores:
            cost['uclass'][1]=cores[mastery]
        unlock=[]
        rankcost=[]
        for c in {'ucarb','ustony','uclass','uunit',}:
            key = '{}{}'.format(c, rank)
            print(key)
            price = rows[0][key]
            print(price)
            if price != '':
                if int(price) > 0:
                    core = cost[c]
                    unlock.append('{0}x {1}'.format(price, core))
        for c in {'rgold','runit',}:
            key = '{}{}'.format(c, rank)
            price = rows[0][key]
            print(price)
            if price != '':
                if int(price) > 0:
                    core = cost[c]
                    rankcost.append('{0}x {1}'.format(price, core))

        return unlock, rankcost

    def _get_text(self, mastery, rank):
        rows = csv_get_rows(self.dataset,'Mastery',mastery)
        for row in rows:
            text.append(row['Text'].format(row[str(rank)]))
        return text

def load_csv(filename):
    return csv.DictReader(open(filename))

def get_csv_row(filecsv, column, match_val, default=None):
    print(match_val)
    csvfile = load_csv(filecsv)
    for row in csvfile:
        if row[column] == match_val:
            if default is not None:
                for k, v in row.items():
                    if v == '':
                        row[k] = default
            return row

def get_csv_rows(filecsv, column, match_val, default=None):
    print(match_val)
    csvfile = load_csv(filecsv)
    package =[]
    for row in csvfile:
        if row[column] == match_val:
            if default is not None:
                for k, v in row.items():
                    if v == '':
                        row[k] = default
            package.append(row)
    return package

def setup(bot):
    bot.add_cog(MCOCTools(bot))
<|endoftext|>