<commit_before># -*- coding: utf-8 -*-
""" controller actions for Contest, Contestand and Founder

    :copyright: Copyright (c) 2014 Bivio Software, Inc.  All Rights Reserved.
    :license: Apache, see LICENSE for more details.
"""

import flask
from functools import wraps
import io
import publicprize.contest.form as pcf
import publicprize.contest.model as pcm
import publicprize.controller as ppc
import publicprize.auth.model as pam
import random
import sqlalchemy.orm
import werkzeug.exceptions

def user_is_admin(func):
    """Require the current user is an administrator."""
    @wraps(func)
    def decorated_function(*args, **kwargs):
        """Forbidden unless allowed."""
        if pam.Admin.is_admin():
            return func(*args, **kwargs)
        werkzeug.exceptions.abort(403)
    return decorated_function

def user_is_contestant_founder_or_admin(func):
    """Require the current user is the contestant founder for an expired
    contest or an administrator."""
    @wraps(func)
    def decorated_function(*args, **kwargs):
        """Forbidden unless allowed."""
        contestant = args[0]
        if pam.Admin.is_admin() or (contestant.is_founder() \
            and contestant.get_contest().is_expired() \
            and contestant.get_contest().is_scoring_completed):
            return func(*args, **kwargs)
        werkzeug.exceptions.abort(403)
    return decorated_function

def user_is_judge(func):
    """Require the current user is a contest judge."""
    @wraps(func)
    def decorated_function(*args, **kwargs):
        """Forbidden unless allowed."""
        if args[0].is_judge():
            return func(*args, **kwargs)
        werkzeug.exceptions.abort(403)
    return decorated_function

class Contest(ppc.Task):
    """Contest actions"""
    def action_about(biv_obj):
        """About page"""
        return Contest._render_template(biv_obj, 'about')

    @ppc.login_required
    @user_is_admin
    def action_admin(biv_obj):
        """Contest administration"""
        return Contest._render_template(biv_obj, 'admin')

    def action_contestants(biv_obj):
        """Public contestant list"""
        return Contest._render_template(
            biv_obj,
            'contestants',
            contest_url=biv_obj.format_absolute_uri(),
        )

    def action_contestants_new(biv_obj):
        return Contest._render_template(biv_obj, 'contestants-new')

    def action_submit_website(biv_obj):
        """Project detail page, loads contest owner and project founders"""
        return pcf.Website().execute(biv_obj)

    def action_submitted_websites(biv_obj):
        """Public list of nominated websites"""
        return Contest._render_template(biv_obj, 'submitted-websites')

    def action_index(biv_obj):
        """Default to contestant list"""
        return Contest.action_contestants(biv_obj)

    def action_judges(biv_obj):
        """List of judges page"""
        return Contest._render_template(biv_obj, 'judges')

    @ppc.login_required
    @user_is_judge
    def action_judging(biv_obj):
        """List of contestants for judgement"""
        return Contest._render_template(
            biv_obj,
            'judging',
        )

    def action_logo(biv_obj):
        """Contestant logo image"""
        return flask.send_file(
            io.BytesIO(biv_obj.contest_logo),
            'image/{}'.format(biv_obj.logo_type)
        )

    def action_new_test_judge(biv_obj):
        """Creates a new test user and judge models and log in."""
        # will raise an exception unless TEST_USER is configured
        flask.g.pub_obj.task_class().action_new_test_user()
        judge = pcm.Judge()
        pcm.db.session.add(judge)
        pcm.db.session.flush()
        pam.db.session.add(pam.BivAccess(
            source_biv_id=flask.session['user.biv_id'],
            target_biv_id=judge.biv_id
        ))
        pam.db.session.add(pam.BivAccess(
            source_biv_id=biv_obj.biv_id,
            target_biv_id=judge.biv_id
        ))
        return flask.redirect('/')

    def action_rules(biv_obj):
        return flask.redirect('/static/pdf/rules.pdf')

    @ppc.login_required
    def action_submit_contestant(biv_obj):
        """Submit project page"""
        return pcf.Contestant().execute(biv_obj)

    def _render_template(biv_obj, name, **kwargs):
        """Render the page, putting the selected menu and contest in env"""
        return flask.render_template(
            'contest/{}.html'.format(name),
            contest=biv_obj,
            selected=name,
            **kwargs
        )

class Contestant(ppc.Task):
    """Contestant actions"""
    def action_contestant(biv_obj):
        """Project detail page, loads contest owner and project founders"""
        if biv_obj.is_public or biv_obj.is_under_review:
            if biv_obj.get_contest().is_expired():
                # TODO(pjm): share return value with Donate form
                return flask.render_template(
                    'contest/detail.html',
                    contestant=biv_obj,
                    contest=biv_obj.get_contest(),
                    contestant_url=biv_obj.format_absolute_uri(),
                    contestant_tweet=biv_obj.display_name
                )
            return pcf.Donate().execute(biv_obj)
        werkzeug.exceptions.abort(404)

    def action_donate_cancel(biv_obj):
        """Return from cancelled payment on paypal site"""
        donor = pcm.Donor.unsafe_load_from_session()
        if donor:
            donor.remove_from_session()
            donor.donor_state = 'canceled'
            ppc.db.session.add(donor)
        form = pcf.Donate()
        form.amount.errors = ['Please resubmit your donation.']
        return form.execute(biv_obj)

    def action_donate_done(biv_obj):
        """Execute the payment after returning from paypal"""
        return pcf.Donate().execute_payment(biv_obj)

    def action_index(biv_obj):
        """Default to contestant page"""
        return Contestant.action_contestant(biv_obj)

    @ppc.login_required
    @user_is_judge
    def action_judging(biv_obj):
        """Contestant judgement"""
        return pcf.Judgement().execute(biv_obj)

    @ppc.login_required
    @user_is_contestant_founder_or_admin
    def action_score(biv_obj):
        """Show the score report for the contestant."""
        summary = {
            'amount_score': 0,
            'amount_raised': 0,
            'judge_score': 0,
            'total_score': 0
        }
        for row in biv_obj.get_contest().get_admin_contestants():
            if row['contestant'].biv_id == biv_obj.biv_id:
                summary = row
                break
        scores_by_judge = biv_obj.get_completed_judge_scores()
        judges = []
        for user_id in scores_by_judge.keys():
            judges.append(
                {
                    'biv_id': user_id,
                    'display_name': pam.User.query.filter_by(
                        biv_id=user_id).one().display_name,
                    'general_comment': pcm.JudgeScore.query.filter_by(
                        judge_biv_id=user_id,
                        question_number=0
                     ).first().judge_comment,
                    'judge_total': biv_obj._score_rows(scores_by_judge[user_id])
                }
            )
        random.Random(biv_obj.display_name).shuffle(judges)
        return flask.render_template(
            'contest/score.html',
            contestant=biv_obj,
            contest=biv_obj.get_contest(),
            judges=judges,
            scores_by_judge=scores_by_judge,
            summary=summary
        )

    def action_thank_you(biv_obj):
        """Show a Thank you page with social media links for contestant."""
        return flask.render_template(
            'contest/thank-you.html',
            contestant=biv_obj,
            contest=biv_obj.get_contest(),
            contestant_url=biv_obj.format_absolute_uri(),
            contestant_tweet="I just backed " + biv_obj.display_name
        )

class Founder(ppc.Task):
    """Founder actions"""
    def action_founder_avatar(biv_obj):
        """Founder avatar image"""
        return flask.send_file(
            io.BytesIO(biv_obj.founder_avatar),
            'image/{}'.format(biv_obj.avatar_type)
        )

class Sponsor(ppc.Task):
    """Sponsor actions"""
    def action_sponsor_logo(biv_obj):
        """Sponsor logo image"""
        return flask.send_file(
            io.BytesIO(biv_obj.sponsor_logo),
            'image/{}'.format(biv_obj.logo_type)
        )
<commit_msg>properly commented action_submit_website<commit_after># -*- coding: utf-8 -*-
""" controller actions for Contest, Contestand and Founder

    :copyright: Copyright (c) 2014 Bivio Software, Inc.  All Rights Reserved.
    :license: Apache, see LICENSE for more details.
"""

import flask
from functools import wraps
import io
import publicprize.contest.form as pcf
import publicprize.contest.model as pcm
import publicprize.controller as ppc
import publicprize.auth.model as pam
import random
import sqlalchemy.orm
import werkzeug.exceptions

def user_is_admin(func):
    """Require the current user is an administrator."""
    @wraps(func)
    def decorated_function(*args, **kwargs):
        """Forbidden unless allowed."""
        if pam.Admin.is_admin():
            return func(*args, **kwargs)
        werkzeug.exceptions.abort(403)
    return decorated_function

def user_is_contestant_founder_or_admin(func):
    """Require the current user is the contestant founder for an expired
    contest or an administrator."""
    @wraps(func)
    def decorated_function(*args, **kwargs):
        """Forbidden unless allowed."""
        contestant = args[0]
        if pam.Admin.is_admin() or (contestant.is_founder() \
            and contestant.get_contest().is_expired() \
            and contestant.get_contest().is_scoring_completed):
            return func(*args, **kwargs)
        werkzeug.exceptions.abort(403)
    return decorated_function

def user_is_judge(func):
    """Require the current user is a contest judge."""
    @wraps(func)
    def decorated_function(*args, **kwargs):
        """Forbidden unless allowed."""
        if args[0].is_judge():
            return func(*args, **kwargs)
        werkzeug.exceptions.abort(403)
    return decorated_function

class Contest(ppc.Task):
    """Contest actions"""
    def action_about(biv_obj):
        """About page"""
        return Contest._render_template(biv_obj, 'about')

    @ppc.login_required
    @user_is_admin
    def action_admin(biv_obj):
        """Contest administration"""
        return Contest._render_template(biv_obj, 'admin')

    def action_contestants(biv_obj):
        """Public contestant list"""
        return Contest._render_template(
            biv_obj,
            'contestants',
            contest_url=biv_obj.format_absolute_uri(),
        )

    def action_contestants_new(biv_obj):
        return Contest._render_template(biv_obj, 'contestants-new')

    def action_submit_website(biv_obj):
        """Page where users can nominate websites to be submitted"""
        return pcf.Website().execute(biv_obj)

    def action_submitted_websites(biv_obj):
        """Public list of nominated websites"""
        return Contest._render_template(biv_obj, 'submitted-websites')

    def action_index(biv_obj):
        """Default to contestant list"""
        return Contest.action_contestants(biv_obj)

    def action_judges(biv_obj):
        """List of judges page"""
        return Contest._render_template(biv_obj, 'judges')

    @ppc.login_required
    @user_is_judge
    def action_judging(biv_obj):
        """List of contestants for judgement"""
        return Contest._render_template(
            biv_obj,
            'judging',
        )

    def action_logo(biv_obj):
        """Contestant logo image"""
        return flask.send_file(
            io.BytesIO(biv_obj.contest_logo),
            'image/{}'.format(biv_obj.logo_type)
        )

    def action_new_test_judge(biv_obj):
        """Creates a new test user and judge models and log in."""
        # will raise an exception unless TEST_USER is configured
        flask.g.pub_obj.task_class().action_new_test_user()
        judge = pcm.Judge()
        pcm.db.session.add(judge)
        pcm.db.session.flush()
        pam.db.session.add(pam.BivAccess(
            source_biv_id=flask.session['user.biv_id'],
            target_biv_id=judge.biv_id
        ))
        pam.db.session.add(pam.BivAccess(
            source_biv_id=biv_obj.biv_id,
            target_biv_id=judge.biv_id
        ))
        return flask.redirect('/')

    def action_rules(biv_obj):
        return flask.redirect('/static/pdf/rules.pdf')

    @ppc.login_required
    def action_submit_contestant(biv_obj):
        """Submit project page"""
        return pcf.Contestant().execute(biv_obj)

    def _render_template(biv_obj, name, **kwargs):
        """Render the page, putting the selected menu and contest in env"""
        return flask.render_template(
            'contest/{}.html'.format(name),
            contest=biv_obj,
            selected=name,
            **kwargs
        )

class Contestant(ppc.Task):
    """Contestant actions"""
    def action_contestant(biv_obj):
        """Project detail page, loads contest owner and project founders"""
        if biv_obj.is_public or biv_obj.is_under_review:
            if biv_obj.get_contest().is_expired():
                # TODO(pjm): share return value with Donate form
                return flask.render_template(
                    'contest/detail.html',
                    contestant=biv_obj,
                    contest=biv_obj.get_contest(),
                    contestant_url=biv_obj.format_absolute_uri(),
                    contestant_tweet=biv_obj.display_name
                )
            return pcf.Donate().execute(biv_obj)
        werkzeug.exceptions.abort(404)

    def action_donate_cancel(biv_obj):
        """Return from cancelled payment on paypal site"""
        donor = pcm.Donor.unsafe_load_from_session()
        if donor:
            donor.remove_from_session()
            donor.donor_state = 'canceled'
            ppc.db.session.add(donor)
        form = pcf.Donate()
        form.amount.errors = ['Please resubmit your donation.']
        return form.execute(biv_obj)

    def action_donate_done(biv_obj):
        """Execute the payment after returning from paypal"""
        return pcf.Donate().execute_payment(biv_obj)

    def action_index(biv_obj):
        """Default to contestant page"""
        return Contestant.action_contestant(biv_obj)

    @ppc.login_required
    @user_is_judge
    def action_judging(biv_obj):
        """Contestant judgement"""
        return pcf.Judgement().execute(biv_obj)

    @ppc.login_required
    @user_is_contestant_founder_or_admin
    def action_score(biv_obj):
        """Show the score report for the contestant."""
        summary = {
            'amount_score': 0,
            'amount_raised': 0,
            'judge_score': 0,
            'total_score': 0
        }
        for row in biv_obj.get_contest().get_admin_contestants():
            if row['contestant'].biv_id == biv_obj.biv_id:
                summary = row
                break
        scores_by_judge = biv_obj.get_completed_judge_scores()
        judges = []
        for user_id in scores_by_judge.keys():
            judges.append(
                {
                    'biv_id': user_id,
                    'display_name': pam.User.query.filter_by(
                        biv_id=user_id).one().display_name,
                    'general_comment': pcm.JudgeScore.query.filter_by(
                        judge_biv_id=user_id,
                        question_number=0
                     ).first().judge_comment,
                    'judge_total': biv_obj._score_rows(scores_by_judge[user_id])
                }
            )
        random.Random(biv_obj.display_name).shuffle(judges)
        return flask.render_template(
            'contest/score.html',
            contestant=biv_obj,
            contest=biv_obj.get_contest(),
            judges=judges,
            scores_by_judge=scores_by_judge,
            summary=summary
        )

    def action_thank_you(biv_obj):
        """Show a Thank you page with social media links for contestant."""
        return flask.render_template(
            'contest/thank-you.html',
            contestant=biv_obj,
            contest=biv_obj.get_contest(),
            contestant_url=biv_obj.format_absolute_uri(),
            contestant_tweet="I just backed " + biv_obj.display_name
        )

class Founder(ppc.Task):
    """Founder actions"""
    def action_founder_avatar(biv_obj):
        """Founder avatar image"""
        return flask.send_file(
            io.BytesIO(biv_obj.founder_avatar),
            'image/{}'.format(biv_obj.avatar_type)
        )

class Sponsor(ppc.Task):
    """Sponsor actions"""
    def action_sponsor_logo(biv_obj):
        """Sponsor logo image"""
        return flask.send_file(
            io.BytesIO(biv_obj.sponsor_logo),
            'image/{}'.format(biv_obj.logo_type)
        )
<|endoftext|>


<commit_before>import asyncio, discord
from   discord.ext import commands
from   Cogs import Utils, DisplayName, Message

def setup(bot):
	# Add the bot
	settings = bot.get_cog("Settings")
	bot.add_cog(Quote(bot, settings))

class Quote(commands.Cog):

	# Init with the bot reference, and a reference to the settings var
	def __init__(self, bot, settings):
		self.bot = bot
		self.settings = settings
		global Utils, DisplayName
		Utils = self.bot.get_cog("Utils")
		DisplayName = self.bot.get_cog("DisplayName")

	@commands.Cog.listener()
	async def on_reaction_add(self, reaction, member):
		# Catch reactions and see if they match our list
		if not type(member) is discord.Member:
			# Not in a server
			return

		# Check if the reaction is added by the bot
		if member.id == self.bot.user.id:
			return

		r =         self.settings.getServerStat(member.guild, "QuoteReaction")
		r_channel = self.settings.getServerStat(member.guild, "QuoteChannel")
		r_admin   = self.settings.getServerStat(member.guild, "QuoteAdminOnly")

		if r == None or r_channel == None:
			# Not setup
			return

		# Check reactions
		if not str(reaction.emoji) == r:
			# Our reaction isn't in there
			return

		em = None
		for reac in reaction.message.reactions:
			if str(reac) == str(r):
				em = reac

		if not em:
			# Broken for no reason?
			return
		ctx = await self.bot.get_context(reaction.message)
		if em.count > 1:
			# Our reaction is already here
			if not r_admin:
				# We're not worried about admin stuffs
				# and someone already quoted
				return
			# Check for admin/bot-admin
			if not Utils.is_bot_admin(ctx,member=member):
				# We ARE worried about admin - and we're not admin... skip
				return
			# Iterate through those that reacted and see if any are admin
			r_users = await reaction.users().flatten()
			for r_user in r_users:
				if r_user == member:
					continue
				if Utils.is_bot_admin(ctx,member=r_user):
					# An admin already quoted - skip
					return
		else:
			# This is the first reaction
			# Check for admin/bot-admin
			if r_admin and not Utils.is_bot_admin(ctx,member=member):
				return

		r_channel = member.guild.get_channel(int(r_channel))
		if r_channel == None:
			# Not a valid channel
			return

		if not len(reaction.message.content) and not len(reaction.message.attachments):
			# Only quote actual text or attachments - no embeds
			return

		# Initialize our message
		msg = reaction.message.content if len(reaction.message.content) else ""

		if len(reaction.message.attachments):
			# We have some attachments to work through
			attach_text = ""
			for a in reaction.message.attachments:
				# Add each attachment by name as a link to its own url
				attach_text += "[{}]({}), ".format(a.filename, a.url)
			# Remove the last ", "
			attach_text = attach_text[:-2]
			msg += "\n\n" + attach_text

		# Build an embed!
		e = {
			"author" : reaction.message.author,
			"pm_after" : -1, # Don't pm quotes
			"description" : msg + "\n\nSent by {} in {} | {} | {} UTC".format(
				reaction.message.author.mention,
				reaction.message.channel.mention,
				"[Link](https://discordapp.com/channels/{}/{}/{})".format(reaction.message.guild.id, reaction.message.channel.id, reaction.message.id),
				reaction.message.created_at.strftime("%I:%M %p")
			),
			"color" : reaction.message.author,
			"footer" : "Quoted by {}#{}".format(member.name, member.discriminator)
		}
		await Message.EmbedText(**e).send(r_channel)

	@commands.command(pass_context=True)
	async def setquotechannel(self, ctx, channel = None):
		"""Sets the channel for quoted messages or disables it if no channel sent (admin only)."""
		if not await Utils.is_bot_admin_reply(ctx): return
		if channel == None:
			self.settings.setServerStat(ctx.guild, "QuoteChannel", None)
			msg = 'Quote channel *disabled*.'
			return await ctx.send(msg)
		channel = DisplayName.channelForName(channel, ctx.guild, "text")
		if channel == None:
			return await ctx.send("I couldn't find that channel :(")
		self.settings.setServerStat(ctx.guild, "QuoteChannel", channel.id)

		msg = 'Quote channel set to {}'.format(channel.mention)
		await ctx.send(msg)

	@commands.command(pass_context=True)
	async def quotechannel(self, ctx):
		"""Prints the current quote channel."""
		qChan = self.settings.getServerStat(ctx.guild, "QuoteChannel")
		if not qChan:
			return await ctx.send("Quoting is currently *disabled*.")
		channel = DisplayName.channelForName(str(qChan), ctx.guild, "text")
		if channel:
			return await ctx.send("The current quote channel is {}".format(channel.mention))
		await ctx.send("Channel id: *{}* no longer exists on this server.  Consider updating this setting!".format(qChan))

	@commands.command(pass_context=True)
	async def clearquotereaction(self, ctx):
		"""Clears the trigger reaction for quoting messages (admin only)."""
		if not await Utils.is_bot_admin_reply(ctx): return
		self.settings.setServerStat(ctx.guild, "QuoteReaction", None)
		await ctx.send("Quote reaction *cleared*.")

	@commands.command(pass_context=True)
	async def setquotereaction(self, ctx):
		"""Sets the trigger reaction for quoting messages (bot-admin only)."""
		if not await Utils.is_bot_admin_reply(ctx): return
		message = await ctx.send("Please react to this message with the desired quote reaction.")
		# Backup then clear - so we don't trigger quoting during this
		backup_reaction = self.settings.getServerStat(ctx.guild, "QuoteReaction")
		self.settings.setServerStat(ctx.guild, "QuoteReaction", None)
		# Now we would wait...
		def check(reaction, user):
			return reaction.message.id == message.id and user == ctx.author
		try:
			reaction, user = await self.bot.wait_for('reaction_add', timeout=60, check=check)
		except:
			# Didn't get a reaction
			self.settings.setServerStat(ctx.guild, "QuoteReaction", backup_reaction)
			return await message.edit(content="Looks like we ran out of time - run `{}setquotereaction` to try again.".format(ctx.prefix))

		# Got it!
		self.settings.setServerStat(ctx.guild, "QuoteReaction", str(reaction.emoji))

		await message.edit(content="Quote reaction set to {}".format(str(reaction.emoji)))

	@commands.command(pass_context=True)
	async def getquotereaction(self, ctx):
		"""Displays the quote reaction if there is one."""
		r = self.settings.getServerStat(ctx.guild, "QuoteReaction")
		await ctx.send("No quote reaction set." if not r else "Current quote reaction is {}".format(r))

	@commands.command(pass_context=True)
	async def quoteadminonly(self, ctx, *, yes_no = None):
		"""Sets whether only admins/bot-admins can quote or not (bot-admin only)."""
		if not await Utils.is_bot_admin_reply(ctx): return
		await ctx.send(Utils.yes_no_setting(ctx,"Admin-only quotes","QuoteAdminOnly",yes_no))
<commit_msg>Add image preview if found in quoted message<commit_after>import asyncio, discord
from   discord.ext import commands
from   Cogs import Utils, DisplayName, Message

def setup(bot):
	# Add the bot
	settings = bot.get_cog("Settings")
	bot.add_cog(Quote(bot, settings))

class Quote(commands.Cog):

	# Init with the bot reference, and a reference to the settings var
	def __init__(self, bot, settings):
		self.bot = bot
		self.settings = settings
		global Utils, DisplayName
		Utils = self.bot.get_cog("Utils")
		DisplayName = self.bot.get_cog("DisplayName")

	@commands.Cog.listener()
	async def on_reaction_add(self, reaction, member):
		# Catch reactions and see if they match our list
		if not type(member) is discord.Member:
			# Not in a server
			return

		# Check if the reaction is added by the bot
		if member.id == self.bot.user.id:
			return

		r =         self.settings.getServerStat(member.guild, "QuoteReaction")
		r_channel = self.settings.getServerStat(member.guild, "QuoteChannel")
		r_admin   = self.settings.getServerStat(member.guild, "QuoteAdminOnly")

		if r == None or r_channel == None:
			# Not setup
			return

		# Check reactions
		if not str(reaction.emoji) == r:
			# Our reaction isn't in there
			return

		em = None
		for reac in reaction.message.reactions:
			if str(reac) == str(r):
				em = reac

		if not em:
			# Broken for no reason?
			return
		ctx = await self.bot.get_context(reaction.message)
		if em.count > 1:
			# Our reaction is already here
			if not r_admin:
				# We're not worried about admin stuffs
				# and someone already quoted
				return
			# Check for admin/bot-admin
			if not Utils.is_bot_admin(ctx,member=member):
				# We ARE worried about admin - and we're not admin... skip
				return
			# Iterate through those that reacted and see if any are admin
			r_users = await reaction.users().flatten()
			for r_user in r_users:
				if r_user == member:
					continue
				if Utils.is_bot_admin(ctx,member=r_user):
					# An admin already quoted - skip
					return
		else:
			# This is the first reaction
			# Check for admin/bot-admin
			if r_admin and not Utils.is_bot_admin(ctx,member=member):
				return

		r_channel = member.guild.get_channel(int(r_channel))
		if r_channel == None:
			# Not a valid channel
			return

		if not len(reaction.message.content) and not len(reaction.message.attachments):
			# Only quote actual text or attachments - no embeds
			return

		# Initialize our message and image field
		msg = reaction.message.content if len(reaction.message.content) else ""
		image = None
		if len(reaction.message.attachments):
			# We have some attachments to work through
			attach_text = ""
			for a in reaction.message.attachments:
				# Add each attachment by name as a link to its own url
				attach_text += "[{}]({}), ".format(a.filename, a.url)
				if image == None and a.filename.lower().endswith((".jpg",".jpeg",".png",".gif")):
					# We got the first image in the attachment list - set it
					image = a.url
			# Remove the last ", "
			attach_text = attach_text[:-2]
			msg += "\n\n" + attach_text
		if len(reaction.message.embeds) and image == None:
			# We have embeds to look at too, and we haven't set an image yet
			for e in reaction.message.embeds:
				if e.thumbnail.url == discord.Embed.Empty: continue
				image = e.thumbnail.url
				break

		# Build an embed!
		e = {
			"author" : reaction.message.author,
			"image" : image,
			"pm_after" : -1, # Don't pm quotes
			"description" : msg + "\n\nSent by {} in {} | {} | {} UTC".format(
				reaction.message.author.mention,
				reaction.message.channel.mention,
				"[Link](https://discord.com/channels/{}/{}/{})".format(reaction.message.guild.id, reaction.message.channel.id, reaction.message.id),
				reaction.message.created_at.strftime("%I:%M %p")
			),
			"color" : reaction.message.author,
			"footer" : "Quoted by {}#{}".format(member.name, member.discriminator)
		}
		await Message.EmbedText(**e).send(r_channel)

	@commands.command(pass_context=True)
	async def setquotechannel(self, ctx, channel = None):
		"""Sets the channel for quoted messages or disables it if no channel sent (admin only)."""
		if not await Utils.is_bot_admin_reply(ctx): return
		if channel == None:
			self.settings.setServerStat(ctx.guild, "QuoteChannel", None)
			msg = 'Quote channel *disabled*.'
			return await ctx.send(msg)
		channel = DisplayName.channelForName(channel, ctx.guild, "text")
		if channel == None:
			return await ctx.send("I couldn't find that channel :(")
		self.settings.setServerStat(ctx.guild, "QuoteChannel", channel.id)

		msg = 'Quote channel set to {}'.format(channel.mention)
		await ctx.send(msg)

	@commands.command(pass_context=True)
	async def quotechannel(self, ctx):
		"""Prints the current quote channel."""
		qChan = self.settings.getServerStat(ctx.guild, "QuoteChannel")
		if not qChan:
			return await ctx.send("Quoting is currently *disabled*.")
		channel = DisplayName.channelForName(str(qChan), ctx.guild, "text")
		if channel:
			return await ctx.send("The current quote channel is {}".format(channel.mention))
		await ctx.send("Channel id: *{}* no longer exists on this server.  Consider updating this setting!".format(qChan))

	@commands.command(pass_context=True)
	async def clearquotereaction(self, ctx):
		"""Clears the trigger reaction for quoting messages (admin only)."""
		if not await Utils.is_bot_admin_reply(ctx): return
		self.settings.setServerStat(ctx.guild, "QuoteReaction", None)
		await ctx.send("Quote reaction *cleared*.")

	@commands.command(pass_context=True)
	async def setquotereaction(self, ctx):
		"""Sets the trigger reaction for quoting messages (bot-admin only)."""
		if not await Utils.is_bot_admin_reply(ctx): return
		message = await ctx.send("Please react to this message with the desired quote reaction.")
		# Backup then clear - so we don't trigger quoting during this
		backup_reaction = self.settings.getServerStat(ctx.guild, "QuoteReaction")
		self.settings.setServerStat(ctx.guild, "QuoteReaction", None)
		# Now we would wait...
		def check(reaction, user):
			return reaction.message.id == message.id and user == ctx.author
		try:
			reaction, user = await self.bot.wait_for('reaction_add', timeout=60, check=check)
		except:
			# Didn't get a reaction
			self.settings.setServerStat(ctx.guild, "QuoteReaction", backup_reaction)
			return await message.edit(content="Looks like we ran out of time - run `{}setquotereaction` to try again.".format(ctx.prefix))

		# Got it!
		self.settings.setServerStat(ctx.guild, "QuoteReaction", str(reaction.emoji))

		await message.edit(content="Quote reaction set to {}".format(str(reaction.emoji)))

	@commands.command(pass_context=True)
	async def getquotereaction(self, ctx):
		"""Displays the quote reaction if there is one."""
		r = self.settings.getServerStat(ctx.guild, "QuoteReaction")
		await ctx.send("No quote reaction set." if not r else "Current quote reaction is {}".format(r))

	@commands.command(pass_context=True)
	async def quoteadminonly(self, ctx, *, yes_no = None):
		"""Sets whether only admins/bot-admins can quote or not (bot-admin only)."""
		if not await Utils.is_bot_admin_reply(ctx): return
		await ctx.send(Utils.yes_no_setting(ctx,"Admin-only quotes","QuoteAdminOnly",yes_no))
<|endoftext|>


<commit_before>import os
BASE_DIR = os.path.dirname(os.path.dirname(__file__))

SECRET_KEY = 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx'

DEBUG = True

INSTALLED_APPS = (
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    'django.contrib.humanize',
    'compressor',
    'templatetag_handlebars',
    'supporttools',
    'userservice',
    'rc_django',
    'grade_conversion_calculator',
    'course_grader.apps.CourseGraderConfig',
)

MIDDLEWARE = (
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.auth.middleware.RemoteUserMiddleware',
    'django.contrib.auth.middleware.SessionAuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
)

AUTHENTICATION_BACKENDS = (
    'django.contrib.auth.backends.RemoteUserBackend',
)

ROOT_URLCONF = 'travis-ci.urls'

WSGI_APPLICATION = 'travis-ci.wsgi.application'

TEMPLATES = [
    {
        'BACKEND': 'django.template.backends.django.DjangoTemplates',
        'APP_DIRS': True,
    },
]

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': os.path.join(BASE_DIR, 'db.sqlite3'),
    }
}

LANGUAGE_CODE = 'en-us'

TIME_ZONE = 'UTC'

USE_I18N = True

USE_L10N = True

USE_TZ = True

STATIC_URL = '/static/'
STATIC_ROOT = BASE_DIR

LOGGING = {
    'version': 1,
    'disable_existing_loggers': False,
    'filters': {
        'add_user': {
            '()': 'course_grader.log.UserFilter'
        },
    },
    'formatters': {
        'course_grader': {
            'format': '%(asctime)s %(user)s %(actas)s %(message)s',
            'datefmt': '[%Y-%m-%d %H:%M:%S]',
        },
    },
    'handlers': {
        'course_grader': {
            'class': 'logging.StreamHandler',
            'filters': ['add_user'],
            'formatter': 'course_grader',
        },
    },
    'loggers': {
        'course_grader': {
            'handlers': ['course_grader'],
            'level': 'DEBUG',
            'propagate': True,
        },
    }
}

GRADEPAGE_ADMIN_GROUP = ''
<commit_msg>drops SessionAuthenticationMiddleware<commit_after>import os
BASE_DIR = os.path.dirname(os.path.dirname(__file__))

SECRET_KEY = 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx'

DEBUG = True

INSTALLED_APPS = (
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    'django.contrib.humanize',
    'compressor',
    'templatetag_handlebars',
    'supporttools',
    'userservice',
    'rc_django',
    'grade_conversion_calculator',
    'course_grader.apps.CourseGraderConfig',
)

MIDDLEWARE = (
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.auth.middleware.RemoteUserMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
)

AUTHENTICATION_BACKENDS = (
    'django.contrib.auth.backends.RemoteUserBackend',
)

ROOT_URLCONF = 'travis-ci.urls'

WSGI_APPLICATION = 'travis-ci.wsgi.application'

TEMPLATES = [
    {
        'BACKEND': 'django.template.backends.django.DjangoTemplates',
        'APP_DIRS': True,
    },
]

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': os.path.join(BASE_DIR, 'db.sqlite3'),
    }
}

LANGUAGE_CODE = 'en-us'

TIME_ZONE = 'UTC'

USE_I18N = True

USE_L10N = True

USE_TZ = True

STATIC_URL = '/static/'
STATIC_ROOT = BASE_DIR

LOGGING = {
    'version': 1,
    'disable_existing_loggers': False,
    'filters': {
        'add_user': {
            '()': 'course_grader.log.UserFilter'
        },
    },
    'formatters': {
        'course_grader': {
            'format': '%(asctime)s %(user)s %(actas)s %(message)s',
            'datefmt': '[%Y-%m-%d %H:%M:%S]',
        },
    },
    'handlers': {
        'course_grader': {
            'class': 'logging.StreamHandler',
            'filters': ['add_user'],
            'formatter': 'course_grader',
        },
    },
    'loggers': {
        'course_grader': {
            'handlers': ['course_grader'],
            'level': 'DEBUG',
            'propagate': True,
        },
    }
}

GRADEPAGE_ADMIN_GROUP = ''
<|endoftext|>


<commit_before>"""Objects for the IL->ASM stage of the compiler."""

import ctypes
import spots
from il_gen import ILCode
from il_gen import ILValue
from spots import Spot

class ASMCode:
    """Stores the ASM code generated from the IL code.

    lines (List) - Lines of ASM code recorded. The commands are stored as
    tuples in this list, where the first value is the name of the command and
    the next values are the command arguments.

    """

    def __init__(self):
        """Initialize ASMCode."""
        self.lines = []

    def add_command(self, command, arg1=None, arg2=None):
        """Add a command to the code.

        command (str) - Name of the command to add.
        arg1 (str) - First argument of the command.
        arg2 (str) - Second argument of the command.

        """
        self.lines.append((command, arg1, arg2))

    def full_code(self):  # noqa: D202
        """Produce the full assembly code.

        return (str) - The assembly code, ready for saving to disk and
        assembling.

        """

        def to_string(line):
            """Convert the provided tuple into a string of asm code.

            Does not terminate with a newline.

            """
            line_str = "     " + line[0]
            if line[1]:
                line_str += " " + line[1]
            if line[2]:
                line_str += ", " + line[2]
            return line_str

        # This code starts every asm program so far, so we put it here.
        header = ["global main", "", "main:"]
        return "\n".join(header + [to_string(line) for line in self.lines])

class ValueMap:
    """Tracks the correspondence between IL values and machine storage.

    value_to_spots (Dict(ILValue -> Set(Spot))) - Map from a IL value to a list
    of spots that IL value is stored at.
    spot_to_values (Dict(Spot -> Set(ILValue))) - Map from a spot to a list of
    IL values that are stored at that spot.

    """

    # List of all registers recognized.
    REGISTERS = [spots.RAX, spots.RSI, spots.RDX]

    def __init__(self):
        """Initialize value map."""
        self.value_to_spots = dict()
        self.spot_to_values = dict()

    def spots(self, il_value):
        """Return all spots from which we can currently get this IL value."""
        if il_value not in self.value_to_spots:
            if il_value.value_type == ILValue.TEMP:
                return set()
            elif il_value.value_type == ILValue.VARIABLE:
                spot = Spot(Spot.STACK, -il_value.offset)
                self._add_value_spot_pair(il_value, spot)
            elif il_value.value_type == ILValue.LITERAL:
                spot = Spot(Spot.LITERAL, il_value.value)
                self._add_value_spot_pair(il_value, spot)

        return self.value_to_spots[il_value]

    def values(self, spot):
        """Return all IL values which are currently stored at this spot."""
        if spot not in self.spot_to_values:
            return set()
        else:
            return self.spot_to_values[spot]

    def set_spots(self, il_value, new_spots):
        """Update the IL value to be stored in the given spots.

        That is, we remove IL value from all spots it was previously stored in,
        and add it to all of the spots provided.

        """
        old_spots = set(self.spots(il_value))
        new_spots = set(new_spots)
        for spot in old_spots:
            self._remove_value_spot_pair(il_value, spot)
        for spot in new_spots:
            self._add_value_spot_pair(il_value, spot)

    def forget(self, il_value):
        """Remove all traces of the IL value from the value map.

        Should be called after the last use of an IL value to free up any
        regsters it may have been holding.

        """
        self.set_spots(il_value, set())

    def get_free_register(self):
        """Return a free register, or None if all registers are occupied."""
        for reg in self.REGISTERS:
            if not self.values(reg):
                return reg

    def _add_value_spot_pair(self, il_value, spot):
        """Add a value-spot pair without considering value/spot types.

        Even internally, use this and the _remove_spot... function instead of
        directly accessing the value_to_spots and spot_to_values dictionaries,
        because these functions make sure the two dictionaries are kept in
        sync.

        """
        if spot not in self.spot_to_values:
            self.spot_to_values[spot] = set()
        if il_value not in self.value_to_spots:
            self.value_to_spots[il_value] = set()

        self.value_to_spots[il_value].add(spot)
        self.spot_to_values[spot].add(il_value)

    def _remove_value_spot_pair(self, il_value, spot):
        """Remove a value-spot pair without considering value/spot types.

        This function raises an exception if the pair is not found.

        Even internally, use this and the _add_spot... function instead of
        directly accessing the value_to_spots and spot_to_values dictionaries,
        because these functions make sure the two dictionaries are kept in
        sync.

        """
        self.value_to_spots[il_value].remove(spot)
        self.spot_to_values[spot].remove(il_value)

        if not self.value_to_spots[il_value]:
            del self.value_to_spots[il_value]
        if not self.spot_to_values[spot]:
            del self.spot_to_values[spot]

class ASMGen:
    """Contains the main logic for generation of the ASM from the IL.

    Note: This class will be refactored once it starts to get unweildy, but
    right now it's fine as a single class.

    il_code (ILCode) - IL code to convert to ASM.
    asm_code (ASMCode) - ASMCode object to populate with ASM.

    value_map (ValueMap) - Internal object storing the current value map.
    value_refs (dict(ILValue->List(int))) - Internal object storing the
    references by line number of each IL value.

    """

    def __init__(self, il_code, asm_code):
        """Initialize ASMGen."""
        self.il_code = il_code
        self.asm_code = asm_code
        self.value_map = ValueMap()

        self.value_refs = self.value_references()

    def value_references(self):
        """Get IL value references by line number.

        This functon returns a dictionary mapping every IL value present in the
        provided IL code to an increasing list of line numbers (relative to the
        IL code) where that IL value is referenced. That is, if:

        n in self.value_references(self)[il_value]

        then at least one of self.il_code[n].arg1, self.il_code[n].arg2, or
        self.il_code[n].output is il_value.

        """
        refs = dict()

        def add_to_refs(il_value, line_num):
            if il_value not in refs:
                refs[il_value] = [line_num]
            else:
                refs[il_value].append(line_num)

        for line_num, line in enumerate(self.il_code):
            add_to_refs(line.output, line_num)
            add_to_refs(line.arg1, line_num)
            add_to_refs(line.arg2, line_num)

        return refs

    def is_forgettable(self, il_value, line_num):
        """Check if il_value is used after line_num.

        When we support multiple blocks, this function will become more
        complex.

        """
        return max(self.value_refs[il_value]) <= line_num

    def make_asm(self):
        """Generate ASM code.

        Uses the ASMCode and ILCode objects passed to the constructor.

        """
        for line_num, line in enumerate(self.il_code):
            # If we're producing a temporary output and this is would be the
            # only use, then skip this line.
            if (line.output and line.output.value_type == ILValue.TEMP and
                    self.is_forgettable(line.output, line_num)):
                continue

            if line.command == ILCode.SET:
                arg_spots = self.value_map.spots(line.arg1)
                self.value_map.set_spots(line.output, arg_spots)

                if self.is_forgettable(line.arg1, line_num):
                    self.value_map.forget(line.arg1)
            elif line.command == ILCode.RETURN:
                self.move_to_spot(
                    self.value_map.spots(line.arg1), spots.RAX,
                    line.arg1.ctype.size)
                self.asm_code.add_command("ret")
            elif line.command == ILCode.ADD:
                # TODO: Clean up this function.

                arg1_spots = set(self.value_map.spots(line.arg1))
                arg2_spots = set(self.value_map.spots(line.arg2))

                # Forget the arguments if possible, so we can reuse those
                # registers.
                if self.is_forgettable(line.arg1, line_num):
                    self.value_map.forget(line.arg1)
                if self.is_forgettable(line.arg2, line_num):
                    self.value_map.forget(line.arg2)

                # If they're both literal ints, add them at compile time.
                if (line.arg1.ctype == ctypes.integer and
                        line.arg2.ctype == ctypes.integer and
                        self.get_literal_spot(arg1_spots) and
                        self.get_literal_spot(arg2_spots)):
                    arg1_literal = self.get_literal_spot(arg1_spots)
                    arg2_literal = self.get_literal_spot(arg2_spots)
                    result_detail = str(
                        int(arg1_literal.detail) + int(arg2_literal.detail))
                    self.value_map.set_spots(
                        line.output, {Spot(Spot.LITERAL, result_detail)})
                    continue

                # Check if arg1 was stored in a register that is now free
                arg1_regs = [spot for spot in arg1_spots
                             if spot.spot_type == Spot.REGISTER and
                             not self.value_map.values(spot)]
                if arg1_regs:
                    self.asm_code.add_command(
                        "add", arg1_regs[0].asm_str(line.arg1.ctype.size),
                        self.best_spot(arg2_spots).asm_str(
                            line.arg2.ctype.size))
                    self.value_map.set_spots(line.output, {arg1_regs[0]})
                    continue

                # Check if arg2 was stored in a register that is now free
                arg2_regs = [spot for spot in arg2_spots
                             if spot.spot_type == Spot.REGISTER and
                             not self.value_map.values(spot)]
                if arg2_regs:
                    self.asm_code.add_command(
                        "add", arg2_regs[0].asm_str(line.arg2.ctype.size),
                        self.best_spot(arg1_spots).asm_str(
                            line.arg1.ctype.size))
                    self.value_map.set_spots(line.output, {arg2_regs[0]})
                    continue

                # Try finding a free register.
                reg = self.value_map.get_free_register()
                if reg:
                    # TODO: issues/13
                    arg1_best = self.best_spot(arg1_spots)
                    arg2_best = self.best_spot(arg2_spots)
                    self.asm_code.add_command(
                        "mov", reg.asm_str(line.arg1.ctype.size),
                        arg1_best.asm_str(line.arg1.ctype.size))
                    self.asm_code.add_command(
                        "add", reg.asm_str(line.arg1.ctype.size),
                        arg2_best.asm_str(line.arg1.ctype.size))
                    self.value_map.set_spots(line.output, {reg})
                else:
                    raise NotImplementedError("No free registers")

            elif line.command == ILCode.MULT:
                raise NotImplementedError
            else:
                raise NotImplementedError

    def get_literal_spot(self, spots):
        """Find a literal spot from the given spots, if possible.

        If any of spots is a literal spot, return it. If not, return None.

        """
        for spot in spots:
            if spot.spot_type == Spot.LITERAL:
                return spot

    def best_spot(self, spots):
        """Pick the best spot from the given list of spots.

        Best is a register spot, next best is a literal spot, worst-case is a
        stack spot. Returns the best spot it finds, or None if spots is empty.

        """
        best_spot = None
        for spot in spots:
            # TODO: Pick the best register, not just the first one
            # encountered.
            if spot.spot_type == Spot.REGISTER:
                return spot
            elif spot.spot_type == Spot.LITERAL:
                best_spot = spot
            elif spot.spot_type == Spot.STACK and not best_spot:
                best_spot = spot
        return best_spot

    def move_to_spot(self, from_spots, to_spot, size):
        """Make asm code to move the given IL value to the given spot.

        This function is currently very dangerous--it overwrites the value of
        the spot without updating the value map. Hence, it should only be used
        at the very end of a block. This will be fixed as needed. If the
        il_value is already in the given register, may not produce any ASM
        code.

        """
        if to_spot.spot_type != Spot.REGISTER:
            raise NotImplementedError(
                "Only register spots currently supported")

        # Check if the value is already in the desired spot
        if any(spot == to_spot for spot in from_spots):
            return

        from_spot = self.best_spot(from_spots)

        self.asm_code.add_command("mov", to_spot.asm_str(size),
                                  from_spot.asm_str(size))
<commit_msg>Add TODO<commit_after>"""Objects for the IL->ASM stage of the compiler."""

import ctypes
import spots
from il_gen import ILCode
from il_gen import ILValue
from spots import Spot

class ASMCode:
    """Stores the ASM code generated from the IL code.

    lines (List) - Lines of ASM code recorded. The commands are stored as
    tuples in this list, where the first value is the name of the command and
    the next values are the command arguments.

    """

    def __init__(self):
        """Initialize ASMCode."""
        self.lines = []

    def add_command(self, command, arg1=None, arg2=None):
        """Add a command to the code.

        command (str) - Name of the command to add.
        arg1 (str) - First argument of the command.
        arg2 (str) - Second argument of the command.

        """
        self.lines.append((command, arg1, arg2))

    def full_code(self):  # noqa: D202
        """Produce the full assembly code.

        return (str) - The assembly code, ready for saving to disk and
        assembling.

        """

        def to_string(line):
            """Convert the provided tuple into a string of asm code.

            Does not terminate with a newline.

            """
            line_str = "     " + line[0]
            if line[1]:
                line_str += " " + line[1]
            if line[2]:
                line_str += ", " + line[2]
            return line_str

        # This code starts every asm program so far, so we put it here.
        header = ["global main", "", "main:"]
        return "\n".join(header + [to_string(line) for line in self.lines])

class ValueMap:
    """Tracks the correspondence between IL values and machine storage.

    value_to_spots (Dict(ILValue -> Set(Spot))) - Map from a IL value to a list
    of spots that IL value is stored at.
    spot_to_values (Dict(Spot -> Set(ILValue))) - Map from a spot to a list of
    IL values that are stored at that spot.

    """

    # List of all registers recognized.
    REGISTERS = [spots.RAX, spots.RSI, spots.RDX]

    def __init__(self):
        """Initialize value map."""
        self.value_to_spots = dict()
        self.spot_to_values = dict()

    def spots(self, il_value):
        """Return all spots from which we can currently get this IL value."""
        if il_value not in self.value_to_spots:
            if il_value.value_type == ILValue.TEMP:
                return set()
            elif il_value.value_type == ILValue.VARIABLE:
                spot = Spot(Spot.STACK, -il_value.offset)
                self._add_value_spot_pair(il_value, spot)
            elif il_value.value_type == ILValue.LITERAL:
                spot = Spot(Spot.LITERAL, il_value.value)
                self._add_value_spot_pair(il_value, spot)

        return self.value_to_spots[il_value]

    def values(self, spot):
        """Return all IL values which are currently stored at this spot."""
        if spot not in self.spot_to_values:
            return set()
        else:
            return self.spot_to_values[spot]

    def set_spots(self, il_value, new_spots):
        """Update the IL value to be stored in the given spots.

        That is, we remove IL value from all spots it was previously stored in,
        and add it to all of the spots provided.

        """
        old_spots = set(self.spots(il_value))
        new_spots = set(new_spots)
        for spot in old_spots:
            self._remove_value_spot_pair(il_value, spot)
        for spot in new_spots:
            self._add_value_spot_pair(il_value, spot)

    def forget(self, il_value):
        """Remove all traces of the IL value from the value map.

        Should be called after the last use of an IL value to free up any
        regsters it may have been holding.

        """
        self.set_spots(il_value, set())

    def get_free_register(self):
        """Return a free register, or None if all registers are occupied."""
        for reg in self.REGISTERS:
            if not self.values(reg):
                return reg

    def _add_value_spot_pair(self, il_value, spot):
        """Add a value-spot pair without considering value/spot types.

        Even internally, use this and the _remove_spot... function instead of
        directly accessing the value_to_spots and spot_to_values dictionaries,
        because these functions make sure the two dictionaries are kept in
        sync.

        """
        if spot not in self.spot_to_values:
            self.spot_to_values[spot] = set()
        if il_value not in self.value_to_spots:
            self.value_to_spots[il_value] = set()

        self.value_to_spots[il_value].add(spot)
        self.spot_to_values[spot].add(il_value)

    def _remove_value_spot_pair(self, il_value, spot):
        """Remove a value-spot pair without considering value/spot types.

        This function raises an exception if the pair is not found.

        Even internally, use this and the _add_spot... function instead of
        directly accessing the value_to_spots and spot_to_values dictionaries,
        because these functions make sure the two dictionaries are kept in
        sync.

        """
        self.value_to_spots[il_value].remove(spot)
        self.spot_to_values[spot].remove(il_value)

        if not self.value_to_spots[il_value]:
            del self.value_to_spots[il_value]
        if not self.spot_to_values[spot]:
            del self.spot_to_values[spot]

class ASMGen:
    """Contains the main logic for generation of the ASM from the IL.

    Note: This class will be refactored once it starts to get unweildy, but
    right now it's fine as a single class.

    il_code (ILCode) - IL code to convert to ASM.
    asm_code (ASMCode) - ASMCode object to populate with ASM.

    value_map (ValueMap) - Internal object storing the current value map.
    value_refs (dict(ILValue->List(int))) - Internal object storing the
    references by line number of each IL value.

    """

    def __init__(self, il_code, asm_code):
        """Initialize ASMGen."""
        self.il_code = il_code
        self.asm_code = asm_code
        self.value_map = ValueMap()

        self.value_refs = self.value_references()

    def value_references(self):
        """Get IL value references by line number.

        This functon returns a dictionary mapping every IL value present in the
        provided IL code to an increasing list of line numbers (relative to the
        IL code) where that IL value is referenced. That is, if:

        n in self.value_references(self)[il_value]

        then at least one of self.il_code[n].arg1, self.il_code[n].arg2, or
        self.il_code[n].output is il_value.

        """
        refs = dict()

        def add_to_refs(il_value, line_num):
            if il_value not in refs:
                refs[il_value] = [line_num]
            else:
                refs[il_value].append(line_num)

        for line_num, line in enumerate(self.il_code):
            add_to_refs(line.output, line_num)
            add_to_refs(line.arg1, line_num)
            add_to_refs(line.arg2, line_num)

        return refs

    def is_forgettable(self, il_value, line_num):
        """Check if il_value is used after line_num.

        When we support multiple blocks, this function will become more
        complex.

        """
        return max(self.value_refs[il_value]) <= line_num

    def make_asm(self):
        """Generate ASM code.

        Uses the ASMCode and ILCode objects passed to the constructor.

        """
        for line_num, line in enumerate(self.il_code):
            # If we're producing a temporary output and this is would be the
            # only use, then skip this line. TODO: Once we have pointer
            # liveliness anaysis, we can also forget variables if we're sure
            # their values will not be used.
            if (line.output and line.output.value_type == ILValue.TEMP and
                    self.is_forgettable(line.output, line_num)):
                continue

            if line.command == ILCode.SET:
                arg_spots = self.value_map.spots(line.arg1)
                self.value_map.set_spots(line.output, arg_spots)

                if self.is_forgettable(line.arg1, line_num):
                    self.value_map.forget(line.arg1)
            elif line.command == ILCode.RETURN:
                self.move_to_spot(
                    self.value_map.spots(line.arg1), spots.RAX,
                    line.arg1.ctype.size)
                self.asm_code.add_command("ret")
            elif line.command == ILCode.ADD:
                # TODO: Clean up this function.

                arg1_spots = set(self.value_map.spots(line.arg1))
                arg2_spots = set(self.value_map.spots(line.arg2))

                # Forget the arguments if possible, so we can reuse those
                # registers.
                if self.is_forgettable(line.arg1, line_num):
                    self.value_map.forget(line.arg1)
                if self.is_forgettable(line.arg2, line_num):
                    self.value_map.forget(line.arg2)

                # If they're both literal ints, add them at compile time.
                if (line.arg1.ctype == ctypes.integer and
                        line.arg2.ctype == ctypes.integer and
                        self.get_literal_spot(arg1_spots) and
                        self.get_literal_spot(arg2_spots)):
                    arg1_literal = self.get_literal_spot(arg1_spots)
                    arg2_literal = self.get_literal_spot(arg2_spots)
                    result_detail = str(
                        int(arg1_literal.detail) + int(arg2_literal.detail))
                    self.value_map.set_spots(
                        line.output, {Spot(Spot.LITERAL, result_detail)})
                    continue

                # Check if arg1 was stored in a register that is now free
                arg1_regs = [spot for spot in arg1_spots
                             if spot.spot_type == Spot.REGISTER and
                             not self.value_map.values(spot)]
                if arg1_regs:
                    self.asm_code.add_command(
                        "add", arg1_regs[0].asm_str(line.arg1.ctype.size),
                        self.best_spot(arg2_spots).asm_str(
                            line.arg2.ctype.size))
                    self.value_map.set_spots(line.output, {arg1_regs[0]})
                    continue

                # Check if arg2 was stored in a register that is now free
                arg2_regs = [spot for spot in arg2_spots
                             if spot.spot_type == Spot.REGISTER and
                             not self.value_map.values(spot)]
                if arg2_regs:
                    self.asm_code.add_command(
                        "add", arg2_regs[0].asm_str(line.arg2.ctype.size),
                        self.best_spot(arg1_spots).asm_str(
                            line.arg1.ctype.size))
                    self.value_map.set_spots(line.output, {arg2_regs[0]})
                    continue

                # Try finding a free register.
                reg = self.value_map.get_free_register()
                if reg:
                    # TODO: issues/13
                    arg1_best = self.best_spot(arg1_spots)
                    arg2_best = self.best_spot(arg2_spots)
                    self.asm_code.add_command(
                        "mov", reg.asm_str(line.arg1.ctype.size),
                        arg1_best.asm_str(line.arg1.ctype.size))
                    self.asm_code.add_command(
                        "add", reg.asm_str(line.arg1.ctype.size),
                        arg2_best.asm_str(line.arg1.ctype.size))
                    self.value_map.set_spots(line.output, {reg})
                else:
                    raise NotImplementedError("No free registers")

            elif line.command == ILCode.MULT:
                raise NotImplementedError
            else:
                raise NotImplementedError

    def get_literal_spot(self, spots):
        """Find a literal spot from the given spots, if possible.

        If any of spots is a literal spot, return it. If not, return None.

        """
        for spot in spots:
            if spot.spot_type == Spot.LITERAL:
                return spot

    def best_spot(self, spots):
        """Pick the best spot from the given list of spots.

        Best is a register spot, next best is a literal spot, worst-case is a
        stack spot. Returns the best spot it finds, or None if spots is empty.

        """
        best_spot = None
        for spot in spots:
            # TODO: Pick the best register, not just the first one
            # encountered.
            if spot.spot_type == Spot.REGISTER:
                return spot
            elif spot.spot_type == Spot.LITERAL:
                best_spot = spot
            elif spot.spot_type == Spot.STACK and not best_spot:
                best_spot = spot
        return best_spot

    def move_to_spot(self, from_spots, to_spot, size):
        """Make asm code to move the given IL value to the given spot.

        This function is currently very dangerous--it overwrites the value of
        the spot without updating the value map. Hence, it should only be used
        at the very end of a block. This will be fixed as needed. If the
        il_value is already in the given register, may not produce any ASM
        code.

        """
        if to_spot.spot_type != Spot.REGISTER:
            raise NotImplementedError(
                "Only register spots currently supported")

        # Check if the value is already in the desired spot
        if any(spot == to_spot for spot in from_spots):
            return

        from_spot = self.best_spot(from_spots)

        self.asm_code.add_command("mov", to_spot.asm_str(size),
                                  from_spot.asm_str(size))
<|endoftext|>


<commit_before>"""Methods to smooth tentative prolongation operators"""

__docformat__ = "restructuredtext en"

from pyamg.utils import approximate_spectral_radius, scale_rows

__all__ = ['jacobi_prolongation_smoother', 'energy_prolongation_smoother']

def jacobi_prolongation_smoother(S, T, omega=4.0/3.0):
    """Jacobi prolongation smoother

    Parameters
    ----------
    S : {csr_matrix, bsr_matrix}
        Sparse NxN matrix used for smoothing.  Typically, A or the
        "filtered matrix" obtained from A by lumping weak connections
        onto the diagonal of A.
    T : {csr_matrix, bsr_matrix}
        Tentative prolongator
    omega : {scalar}
        Damping parameter

    Returns
    -------
    P : {csr_matrix, bsr_matrix}
        Smoothed (final) prolongator defined by P = (I - omega/rho(S) S) * T
        where rho(S) is an approximation to the spectral radius of S.

    """

    D = S.diagonal()
    D_inv = 1.0 / D
    D_inv[D == 0] = 0

    D_inv_S = scale_rows(S, D_inv, copy=True)
    D_inv_S *= omega/approximate_spectral_radius(D_inv_S)

    P = T - (D_inv_S*T)

    return P

""" sa_energy_min + helper functions minimize the energy of a tentative prolongator for use in SA """

from numpy import array, zeros, matrix, mat
from scipy.sparse import csr_matrix, isspmatrix_csr, bsr_matrix, isspmatrix_bsr, spdiags
from scipy.linalg import svd, norm
import pyamg
from pyamg.utils import UnAmal, BSR_Get_Colindices, BSR_Get_Row
from scipy.io import loadmat, savemat

########################################################################################################
#   Helper function for the energy minimization prolongator generation routine

def Satisfy_Constraints_CSR(U, Sparsity_Pattern, B, BtBinv):
    """Update U to satisfy U*B = 0

    Input
    =====
    U                       CSR Matrix to operate on
    Sparsity_Pattern        Sparsity pattern to enforce
    B                       Near nullspace vectors
    BtBinv                  Local inv(B'*B) matrices for each dof, i.

    Output
    ======
    Updated U, so that U*B = 0.  Update is computed by orthogonal (in 2-norm)
    projecting out the components of span(B) in U in a row-wise fashion

    """

    Nfine = U.shape[0]

    UB = U*mat(B)
    #Project out U's components in span(B) row-wise
    for i in range(Nfine):
        rowstart = Sparsity_Pattern.indptr[i]
        rowend = Sparsity_Pattern.indptr[i+1]
        length = rowend - rowstart
        colindx = Sparsity_Pattern.indices[rowstart:rowend]

        if(length != 0):
            Bi = B[colindx,:]
            UBi = UB[i,:]

            update_local = (Bi*(BtBinv[i]*UBi.T)).T
            Sparsity_Pattern.data[rowstart:rowend] = update_local

    #Now add in changes from Sparsity_Pattern to U.  We don't write U directly in the above loop,
    #   because its possible, once in a blue moon, to have the sparsity pattern of U be a subset
    #   of the pattern for Sparsity_Pattern.  This is more expensive, but more robust.
    U = U - Sparsity_Pattern
    Sparsity_Pattern.data[:] = 1.0
    return U

def Satisfy_Constraints_BSR(U, Sparsity_Pattern, B, BtBinv, colindices):
    """Update U to satisfy U*B = 0

    Input
    =====
    U                     BSR Matrix to operate on
    Sparsity_Pattern      Sparsity pattern to enforce
    B                     Near nullspace vectors
    BtBinv                Local inv(B'*B) matrices for each dof, i.
    colindices            List indexed by node that returns column indices for
                          all dof's in that node.  Assumes that each block is
                          perfectly dense.  The below code does assure this
                          perfect denseness.

    Output
    ======
    Updated U, so that U*B = 0.  Update is computed by orthogonal (in 2-norm)
    projecting out the components of span(B) in U in a row-wise fashion

    """

    Nfine = U.shape[0]
    RowsPerBlock = U.blocksize[0]
    ColsPerBlock = U.blocksize[1]
    Nnodes = Nfine/RowsPerBlock

    UB = U*mat(B)

    rowoffset = 0
    for i in range(Nnodes):
        rowstart = Sparsity_Pattern.indptr[i]
        rowend = Sparsity_Pattern.indptr[i+1]
        colindx = colindices[i]
        length = len(colindx)
        numBlocks = rowend-rowstart

        if(length != 0):
            Bi = B[colindx,:]
            UBi = UB[rowoffset:(rowoffset+RowsPerBlock), :]
            update_local = (Bi*(BtBinv[i]*UBi.T))

            #Write node's values
            for j in range(RowsPerBlock):
                Sparsity_Pattern.data[rowstart:rowend, j, :] = update_local[:,j].reshape(numBlocks, ColsPerBlock)

        rowoffset += RowsPerBlock

    #Now add in changes from Sparsity_Pattern to U.  We don't write U directly in the above loop,
    #   because its possible, once in a blue moon, to have the sparsity pattern of U be a subset
    #   of the pattern for Sparsity_Pattern.  This is more expensive, but more robust.
    U = U - Sparsity_Pattern
    Sparsity_Pattern.data[:,:,:] = 1.0
    return U

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

def energy_prolongation_smoother(A, T, Atilde, B, SPD=True, num_iters=4, min_tol=1e-8, file_output=False):
    """Minimize the energy of the coarse basis functions (columns of T)

    Parameters
    ----------
    A : {csr_matrix, bsr_matrix}
        Sparse NxN matrix
    T : {bsr_matrix}
        Tentative prolongator, a NxM sparse matrix (M < N)
    Atilde : {csr_matrix}
        Strength of connection matrix
    B : {array}
        Near-nullspace modes for coarse grid.  Has shape (M,k) where
        k is the number of coarse candidate vectors.
    SPD : boolean
        Booolean denoting symmetric positive-definiteness of A
    num_iters : integer
        Number of energy minimization steps to apply to the prolongator
    min_tol : scalar
        Minimization tolerance
    file_output : boolean
        Optional diagnostic file output of matrices

    Returns
    -------
    P : {bsr_matrix}
        Smoothed prolongator

    References
    ----------

        Mandel
        "Energy Optimization of Algebraic Multigrid Bases."

        TODO COMPLETE REF

    """

    #====================================================================
    #Test Inputs
    if( not(isinstance(num_iters,int)) ):
        raise TypeError("\nCalling sa_energy_min Incorrectly.  "
                "Number of minimization iterations applied to P is \"num_iters\", where \"num_iters\" must be an integer.\n\n")
    if( num_iters < 0 ):
        raise TypeError("\nCalling sa_energy_min Incorrectly.  ",
                "Number of minimization iterations applied to P is \"num_iters\", where \"num_iters\" >= 0.\n\n")
    if( min_tol > 1  ):
        raise TypeError("\nCalling sa_energy_min Incorrectly.  0 < \"min_tol\" < 1.\n\n")

    csrflag = isspmatrix_csr(A)
    if( not(csrflag) and (isspmatrix_bsr(A) == False)):
        raise TypeError("sa_energy_min routine requires a CSR or BSR operator.  Aborting.\n")
    if(isspmatrix_bsr(T) == False):
        raise TypeError("sa_energy_min routine requires a BSR tentative prolongator.  Aborting.\n")
    if( not(csrflag) and (T.blocksize[0] != A.blocksize[0]) ):
        print "Warning, T's row-blocksize should be the same as A's blocksize.\n"
    if( (T.nnz == 0) or (Atilde.nnz == 0) or (A.nnz == 0) ):
        print "Error in sa_energy_min(..).  A, T or Atilde has no nonzeros on a level."
        return T

    #====================================================================

    #====================================================================
    # Retrieve problem information
    Nfine = T.shape[0]
    Ncoarse = T.shape[1]
    NullDim = B.shape[1]
    if(csrflag):
        numPDEs = 1
    else:
        #Number of PDEs per point is defined implicitly by block size
        numPDEs = A.blocksize[0]
    #====================================================================

    #====================================================================
    # Construct Dinv and Unamalgate Atilde if (numPDEs > 1)
    D = A.diagonal();
    #Must Do extensive checking for 0 rows of A.
    if(D.nonzero()[0].shape[0] != A.shape[0]):
        zero_rows = D.__eq__(0.0).nonzero()[0]
        for i in zero_rows:
            #See if row i is all zero
            if(csrflag):
                if(A.indptr[i] != A.indptr[i+1]):
                    raise ValueError("Zero on diag(A) for nonzero row of A in sa_ode_energy_min routine. -- Aborting.\n\n")
            else:
                zi, colindx = BSR_Get_Row(A, i)
                if(min(colindx.shape) > 0):
                    raise ValueError("Zero on diag(A) for nonzero row of A in sa_ode_energy_min routine. -- Aborting.\n\n")
        # Zeros on D represent 0 rows, so we can just set D to 1.0 at those locations and then Dinv*A
        #   at the zero rows of A will still be zero
        D[zero_rows] = 1.0

    if(csrflag):
        Dinv    = spdiags( [1.0/D], [0], Nfine, Nfine, format = 'csr')
    else:
        Dinv    = spdiags( [1.0/D], [0], Nfine, Nfine, format = 'bsr')
        Dinv = bsr_matrix(Dinv, blocksize=A.blocksize)

    #If A is BSR, then Atilde needs to be "UnAmalgamated" to generate prolongator
    #   sparsity pattern.  Sparsity pattern is generated through matrix multiplication
    if(not(csrflag)):
        #UnAmal returns a BSR matrix, so the mat-mat will be a BSR mat-mat.  Unfortunately,
        #   we also need column indices for Sparsity_Pattern
        Sparsity_Pattern = UnAmal(Atilde, numPDEs, numPDEs).__abs__()*T.__abs__()
        Sparsity_Pattern.data[:,:,:] = 1.0
        Sparsity_Pattern.sort_indices()
        colindices = BSR_Get_Colindices(Sparsity_Pattern)
    else:
        #Sparsity_Pattern will be CSR as Atilde is CSR.  This means T will be converted to
        #   CSR, but we need Sparsity_Pattern in CSR.  T is converted back to BSR at the end.
        blocks = T.blocksize
        T = T.tocsr()
        Sparsity_Pattern = Atilde.__abs__()*T.__abs__()
        Sparsity_Pattern.data[:] = 1.0
        Sparsity_Pattern.sort_indices()
    #====================================================================

    #====================================================================
    #Optional file output for diagnostic purposes
    if(file_output == True):
        savemat('Sparsity_Pattern', { 'Sparsity_Pattern' : Sparsity_Pattern.toarray() } )
        savemat('Amat', { 'Amat' : A.toarray() } )
        savemat('Atilde', { 'Atilde' : Atilde.toarray() } )
        savemat('P', { 'P' : T.toarray() } )
        savemat('ParamsEnMin', {'nPDE' : numPDEs, 'Nits' : num_iters, 'SPD' : SPD } )
        savemat('Bone', { 'Bone' : array(B) } )
    #====================================================================

    #====================================================================
    #Construct array of inv(Bi'Bi), where Bi is B restricted to row i's sparsity pattern in
    #   Sparsity Pattern.  This array is used multiple times in the Satisfy_Constraints routine.
    if(csrflag):
        preall = zeros((NullDim,NullDim))
        BtBinv = [matrix(preall,copy=True) for i in range(Nfine)]
        del preall
        B = mat(B)
        for i in range(Nfine):
            rowstart = Sparsity_Pattern.indptr[i]
            rowend = Sparsity_Pattern.indptr[i+1]
            length = rowend - rowstart
            colindx = Sparsity_Pattern.indices[rowstart:rowend]

            if(length != 0):
                Bi = B[colindx,:]
                #Calculate SVD as system may be singular
                U,Sigma,VT = svd(Bi.T*Bi)

                #Filter Sigma and calculate inv(Sigma)
                if(abs(Sigma[0]) < 1e-10):
                    Sigma[:] = 0.0
                else:
                    #Zero out "numerically" zero singular values
                    #   Efficiency TODO -- would this be faster in a loop that starts from the
                    #   back of Sigma and assumes Sigma is sorted?  Experiments say no.
                    Sigma =  (Sigma/Sigma[0]).__abs__().__gt__(1e-8)*Sigma

                    #Test for any zeros in Sigma
                    if(Sigma[NullDim-1] == 0.0):
                        #Truncate U, VT and Sigma w.r.t. zero entries in Sigma
                        indys = Sigma.nonzero()[0]
                        Sigma = Sigma[indys]
                        U = U[:,indys]
                        VT = VT[indys,:]

                    #Invert nonzero sing values
                    Sigma = 1.0/Sigma

                #Calculate inverse
                BtBinv[i] = mat(VT.T)*mat(Sigma.reshape(Sigma.shape[0],1)*U.T)
    else:   #BSR matrix
        preall = zeros((NullDim,NullDim))
        RowsPerBlock = Sparsity_Pattern.blocksize[0]
        Nnodes = Nfine/RowsPerBlock
        BtBinv = [matrix(preall,copy=True) for i in range(Nnodes)]
        del preall
        B = mat(B)
        for i in range(Nnodes):

            rowstart = Sparsity_Pattern.indptr[i]
            rowend = Sparsity_Pattern.indptr[i+1]
            length = rowend - rowstart
            colindx = colindices[i]

            if(length != 0):
                Bi = B[colindx,:]
                #Calculate SVD as system may be singular
                U,Sigma,VT = svd(Bi.T*Bi)

                #Filter Sigma and calculate inv(Sigma)
                if(abs(Sigma[0]) < 1e-10):
                    Sigma[:] = 0.0
                else:
                    #Zero out "numerically" zero singular values
                    #   Efficiency TODO -- would this be faster in a loop that starts from the
                    #   back of Sigma and assumes Sigma is sorted?  Experiments say no.
                    Sigma =  (Sigma/Sigma[0]).__abs__().__gt__(1e-8)*Sigma

                    #Test for any zeros in Sigma
                    if(Sigma[NullDim-1] == 0.0):
                        #Truncate U, VT and Sigma w.r.t. zero entries in Sigma
                        indys = Sigma.nonzero()[0]
                        Sigma = Sigma[indys]
                        U = U[:,indys]
                        VT = VT[indys,:]

                    #Invert nonzero sing values
                    Sigma = 1.0/Sigma

                #Calculate inverse
                BtBinv[i] = mat(VT.T)*mat(Sigma.reshape(Sigma.shape[0],1)*U.T)

    #====================================================================

    #====================================================================
    #Prepare for Energy Minimization
    #Calculate initial residual
    R = -A*T

    #Enforce constraints on R.  First the sparsity pattern, then the nullspace vectors.
    R = R.multiply(Sparsity_Pattern)
    if(csrflag):
        R = Satisfy_Constraints_CSR(R, Sparsity_Pattern, B, BtBinv)
    else:
        R = Satisfy_Constraints_BSR(R, Sparsity_Pattern, B, BtBinv, colindices)
    if(R.nnz == 0 ):
        print "Error in sa_energy_min(..).  Initial R no nonzeros on a level.  Calling Default Prolongator Smoother\n"
        T = pyamg.sa.sa_smoothed_prolongator(A,T)
        return T

    #Calculate max norm of the residual
    resid = max(R.data.flatten().__abs__())
    #print "Energy Minimization of Prolongator --- Iteration 0 --- r = " + str(resid)
    #====================================================================

    #====================================================================
    #Iteratively minimize the energy of T subject to the constraints of Sparsity_Pattern
    #   and maintaining T's effect on B, i.e. T*B = (T+Update)*B, i.e. Update*B = 0
    i = 0
    if(SPD):
        #Apply CG
        while( (i < num_iters) and (resid > min_tol) ):
            #Apply diagonal preconditioner
            Z = Dinv*R

            #Frobenius innerproduct of (R,Z) = sum(rk.*zk)
            newsum = (R.multiply(Z)).sum()

            #P is the search direction, not the prolongator, which is T.
            if(i == 0):
                P = Z
            else:
                beta = newsum/oldsum
                P = Z + beta*P
            oldsum = newsum

            #Calculate new direction and enforce constraints
            AP = A*P
            AP = AP.multiply(Sparsity_Pattern)
            if(csrflag):
                AP = Satisfy_Constraints_CSR(AP, Sparsity_Pattern, B, BtBinv)
            else:
                AP = Satisfy_Constraints_BSR(AP, Sparsity_Pattern, B, BtBinv, colindices)

            #Frobenius innerproduct of (P, AP)
            alpha = newsum/(P.multiply(AP)).sum()

            #Update the prolongator, T
            T = T + alpha*P

            #Update residual
            R = R - alpha*AP

            i += 1
            resid = max(R.data.flatten().__abs__())
            #print "Energy Minimization of Prolongator --- Iteration " + str(i) + " --- r = " + str(resid)

    else:
        #Apply min-res to the nonsymmetric system
        while( (i < num_iters) and (resid > min_tol) ):

            #P is the search direction, not the prolongator
            P = A*R

            #Enforce constraints on P
            P = P.multiply(Sparsity_Pattern)
            if(csrflag):
                P = Satisfy_Constraints_CSR(P, Sparsity_Pattern, B, BtBinv)
            else:
                P = Satisfy_Constraints_BSR(P, Sparsity_Pattern, B, BtBinv, colindices)

            #Frobenius innerproduct of (P, R)
            numer = (P.multiply(R)).sum()

            #Frobenius innerproduct of (P, P)
            denom = (P.multiply(P)).sum()

            alpha = numer/denom

            #Update prolongator
            T = T + alpha*R

            #Update residual
            R = R - alpha*P

            i += 1
            resid = max(R.data.flatten().__abs__())
            #print "Energy Minimization of Prolongator --- Iteration " + str(i) + " --- r = " + str(resid)
    #====================================================================

    #====================================================================
    #Make sure T is in correct block format.
    if(csrflag):
        T = T.tobsr(blocksize=blocks)
    #====================================================================

    #====================================================================
    #Optional file output for diagnostic purposes
    if(file_output == True):
        savemat('Ppyth', { 'Ppyth' : T.toarray() } )
    #====================================================================

    return T

<commit_msg>sped up energy min <commit_after>"""Methods to smooth tentative prolongation operators"""

__docformat__ = "restructuredtext en"

from pyamg.utils import approximate_spectral_radius, scale_rows

__all__ = ['jacobi_prolongation_smoother', 'energy_prolongation_smoother']

def jacobi_prolongation_smoother(S, T, omega=4.0/3.0):
    """Jacobi prolongation smoother

    Parameters
    ----------
    S : {csr_matrix, bsr_matrix}
        Sparse NxN matrix used for smoothing.  Typically, A or the
        "filtered matrix" obtained from A by lumping weak connections
        onto the diagonal of A.
    T : {csr_matrix, bsr_matrix}
        Tentative prolongator
    omega : {scalar}
        Damping parameter

    Returns
    -------
    P : {csr_matrix, bsr_matrix}
        Smoothed (final) prolongator defined by P = (I - omega/rho(S) S) * T
        where rho(S) is an approximation to the spectral radius of S.

    """

    D = S.diagonal()
    D_inv = 1.0 / D
    D_inv[D == 0] = 0

    D_inv_S = scale_rows(S, D_inv, copy=True)
    D_inv_S *= omega/approximate_spectral_radius(D_inv_S)

    P = T - (D_inv_S*T)

    return P

""" sa_energy_min + helper functions minimize the energy of a tentative prolongator for use in SA """

from numpy import array, zeros, matrix, mat
from scipy.sparse import csr_matrix, isspmatrix_csr, bsr_matrix, isspmatrix_bsr, spdiags
from scipy.linalg import svd, norm, pinv2
import pyamg
from pyamg.utils import UnAmal, BSR_Get_Colindices, BSR_Get_Row
from scipy.io import loadmat, savemat

########################################################################################################
#   Helper function for the energy minimization prolongator generation routine

def Satisfy_Constraints(U, Sparsity_Pattern, B, BtBinv, colindices):
    """Update U to satisfy U*B = 0

    Input
    =====
    U                     BSR Matrix to operate on
    Sparsity_Pattern      Sparsity pattern to enforce
    B                     Near nullspace vectors
    BtBinv                Local inv(B'*B) matrices for each dof, i.
    colindices            List indexed by node that returns column indices for
                          all dof's in that node.  Assumes that each block is
                          perfectly dense.  The below code does assure this
                          perfect denseness.

    Output
    ======
    Updated U, so that U*B = 0.  Update is computed by orthogonal (in 2-norm)
    projecting out the components of span(B) in U in a row-wise fashion

    """

    Nfine = U.shape[0]
    RowsPerBlock = U.blocksize[0]
    ColsPerBlock = U.blocksize[1]
    Nnodes = Nfine/RowsPerBlock

    UB = U*mat(B)

    from numpy import asarray, dot
    B = asarray(B)
    UB = asarray(UB)
    rows = csr_matrix((U.indices,U.indices,U.indptr), shape=(U.shape[0]/RowsPerBlock,U.shape[1]/ColsPerBlock)).tocoo().row
    B  = B.reshape(-1,ColsPerBlock,B.shape[1])
    UB = UB.reshape(-1,RowsPerBlock,ColsPerBlock)
    for n,j in enumerate(U.indices):
        i = rows[n]
        U.data[n] -= dot(dot(UB[i],BtBinv[i]),B[j])

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

def energy_prolongation_smoother(A, T, Atilde, B, SPD=True, num_iters=4, min_tol=1e-8, file_output=False):
    """Minimize the energy of the coarse basis functions (columns of T)

    Parameters
    ----------
    A : {csr_matrix, bsr_matrix}
        Sparse NxN matrix
    T : {bsr_matrix}
        Tentative prolongator, a NxM sparse matrix (M < N)
    Atilde : {csr_matrix}
        Strength of connection matrix
    B : {array}
        Near-nullspace modes for coarse grid.  Has shape (M,k) where
        k is the number of coarse candidate vectors.
    SPD : boolean
        Booolean denoting symmetric positive-definiteness of A
    num_iters : integer
        Number of energy minimization steps to apply to the prolongator
    min_tol : scalar
        Minimization tolerance
    file_output : boolean
        Optional diagnostic file output of matrices

    Returns
    -------
    P : {bsr_matrix}
        Smoothed prolongator

    References
    ----------

        Mandel
        "Energy Optimization of Algebraic Multigrid Bases."

        TODO COMPLETE REF

    """

    #====================================================================
    #Test Inputs
    if num_iters < 0:
        raise ValueError('num_iters must be > 0')
    if min_tol > 1:
        raise ValueError('min_tol must be <= 1')

    if isspmatrix_csr(A):
        A = A.tobsr(blocksize=(1,1), copy=False)
    elif isspmatrix_bsr(A):
        pass
    else:
        raise TypeError("A must be csr_matrix or bsr_matrix")

    if not isspmatrix_csr(Atilde):
        raise TypeError("Atilde must be csr_matrix")

    csrflag = isspmatrix_csr(A)

    if T.blocksize[0] != A.blocksize[0]:
        raise ValueError("T's row-blocksize should be the same as A's blocksize")

    if min(T.nnz, Atilde.nnz, A.nnz) == 0:
        return T

    #====================================================================

    #====================================================================
    # Retrieve problem information
    Nfine = T.shape[0]
    Ncoarse = T.shape[1]
    NullDim = B.shape[1]
    if(csrflag):
        numPDEs = 1
    else:
        #Number of PDEs per point is defined implicitly by block size
        numPDEs = A.blocksize[0]
    #====================================================================

    #====================================================================
    # Construct Dinv and Unamalgate Atilde if (numPDEs > 1)
    D = A.diagonal();
    if (D == 0).any():
        zero_rows = (D == 0).nonzeros()[0]
        if (diff(A.tocsr().indptr)[zero_rows] > 0).any():
            raise ValueError('zero on diag(A) for nonzero row of A')
        # Zeros on D represent 0 rows, so we can just set D to 1.0 at those locations and then Dinv*A
        #   at the zero rows of A will still be zero
        D[zero_rows] = 1.0
    Dinv = 1.0/D

    #####UnAmal returns a BSR matrix, so the mat-mat will be a BSR mat-mat.  Unfortunately,
    #####   we also need column indices for Sparsity_Pattern
    Sparsity_Pattern = UnAmal(abs(Atilde), numPDEs, numPDEs)*abs(T)
    Sparsity_Pattern.data[:,:,:] = 1.0
    Sparsity_Pattern.sort_indices()
    colindices = BSR_Get_Colindices(Sparsity_Pattern)

    ##TODO remove workaround
    #from numpy import array_split
    #temp = Sparsity_Pattern.tocsr()
    #temp.sort_indices()
    #colindices = array_split(temp.indices, temp.indptr[1:-1])

    ##import pdb; pdb.set_trace()

    #====================================================================
    #Optional file output for diagnostic purposes
    if(file_output == True):
        savemat('Sparsity_Pattern', { 'Sparsity_Pattern' : Sparsity_Pattern.toarray() } )
        savemat('Amat', { 'Amat' : A.toarray() } )
        savemat('Atilde', { 'Atilde' : Atilde.toarray() } )
        savemat('P', { 'P' : T.toarray() } )
        savemat('ParamsEnMin', {'nPDE' : numPDEs, 'Nits' : num_iters, 'SPD' : SPD } )
        savemat('Bone', { 'Bone' : array(B) } )
    #====================================================================

    #====================================================================
    #Construct array of inv(Bi'Bi), where Bi is B restricted to row i's sparsity pattern in
    #   Sparsity Pattern.  This array is used multiple times in the Satisfy_Constraints routine.

    preall = zeros((NullDim,NullDim))
    RowsPerBlock = Sparsity_Pattern.blocksize[0]
    Nnodes = Nfine/RowsPerBlock
    BtBinv = [matrix(preall,copy=True) for i in range(Nnodes)]
    del preall
    B = mat(B)
    for i in range(Nnodes):
        colindx = colindices[i]
        if len(colindx) > 0:
            Bi = B[colindx,:]
            BtBinv[i] = pinv2(Bi.T*Bi)
    #====================================================================

    #====================================================================
    #Prepare for Energy Minimization
    #Calculate initial residual
    R = -A*T

    #Enforce constraints on R.  First the sparsity pattern, then the nullspace vectors.
    R = R.multiply(Sparsity_Pattern)
    Satisfy_Constraints(R, Sparsity_Pattern, B, BtBinv, colindices)

    if R.nnz == 0:
        print "Error in sa_energy_min(..).  Initial R no nonzeros on a level.  Calling Default Prolongator Smoother\n"
        return jacobi_prolongation_smoother(Atilde, T)

    #Calculate max norm of the residual
    resid = abs(R.data.flatten()).max()
    #print "Energy Minimization of Prolongator --- Iteration 0 --- r = " + str(resid)
    #====================================================================

    #====================================================================
    #Iteratively minimize the energy of T subject to the constraints of Sparsity_Pattern
    #   and maintaining T's effect on B, i.e. T*B = (T+Update)*B, i.e. Update*B = 0
    i = 0
    if SPD:
        #Apply CG
        while i < num_iters and resid > min_tol:
            #Apply diagonal preconditioner
            Z = scale_rows(R,Dinv)

            #Frobenius innerproduct of (R,Z) = sum(rk.*zk)
            newsum = (R.multiply(Z)).sum()

            #P is the search direction, not the prolongator, which is T.
            if(i == 0):
                P = Z
            else:
                beta = newsum/oldsum
                P = Z + beta*P
            oldsum = newsum

            #Calculate new direction and enforce constraints
            AP = A*P
            AP = AP.multiply(Sparsity_Pattern)
            Satisfy_Constraints(AP, Sparsity_Pattern, B, BtBinv, colindices)

            #Frobenius innerproduct of (P, AP)
            alpha = newsum/(P.multiply(AP)).sum()

            #Update the prolongator, T
            T = T + alpha*P

            #Update residual
            R = R - alpha*AP

            i += 1
            resid = abs(R.data).max()
            #print "Energy Minimization of Prolongator --- Iteration " + str(i) + " --- r = " + str(resid)

    else:
        #Apply min-res to the nonsymmetric system
        while( (i < num_iters) and (resid > min_tol) ):

            #P is the search direction, not the prolongator
            P = A*R

            #Enforce constraints on P
            P = P.multiply(Sparsity_Pattern)
            if(csrflag):
                P = Satisfy_Constraints_CSR(P, Sparsity_Pattern, B, BtBinv)
            else:
                P = Satisfy_Constraints_BSR(P, Sparsity_Pattern, B, BtBinv, colindices)

            #Frobenius innerproduct of (P, R)
            numer = (P.multiply(R)).sum()

            #Frobenius innerproduct of (P, P)
            denom = (P.multiply(P)).sum()

            alpha = numer/denom

            #Update prolongator
            T = T + alpha*R

            #Update residual
            R = R - alpha*P

            i += 1
            resid = max(R.data.flatten().__abs__())
            #print "Energy Minimization of Prolongator --- Iteration " + str(i) + " --- r = " + str(resid)
    #====================================================================

    #====================================================================
    #Make sure T is in correct block format.
    if(csrflag):
        T = T.tobsr(blocksize=blocks)
    #====================================================================

    #====================================================================
    #Optional file output for diagnostic purposes
    if(file_output == True):
        savemat('Ppyth', { 'Ppyth' : T.toarray() } )
    #====================================================================

    return T

<|endoftext|>