<commit_before>/*
Copyright (c) 2016, Valeriy Soldatov
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
 * Neither the name of the football.mojgorod.ru nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <QCloseEvent>
#include <QDesktopServices>
#include <QDesktopWidget>
#include <QFile>
#include <QList>
#include <QMap>
#include <QMessageBox>
#include <QSettings>
#include <QString>
#include <QTextStream>

#include "MainWindow.h"
#include "TestItem.h"

#include "ui_MainWindow.h"

using namespace std;

MainWindow::MainWindow() {

    widget.setupUi(this);

    connect(widget.actionExit, SIGNAL(triggered()),
            this, SLOT(close()));
    connect(widget.actionAbout, SIGNAL(triggered()),
            this, SLOT(about()));
    connect(widget.btnRefresh, SIGNAL(clicked()),
            this, SLOT(refresh()));
    connect(widget.table, SIGNAL(cellClicked(int, int)),
            this, SLOT(cellClicked(int, int)));
    connect(widget.text, SIGNAL(anchorClicked(const QUrl &)),
            this, SLOT(linkActivated(const QUrl &)));

    QRect rectangle = frameGeometry();
    rectangle.moveCenter(QDesktopWidget().availableGeometry().center());
    move(rectangle.topLeft());

    QSettings settings(QSettings::IniFormat, QSettings::UserScope,
            "QHudsonResults", "QHudsonResults");
    data = settings.value("data").toString();
    if (data == "") {
        QString applicationDir = QCoreApplication::applicationDirPath();
        if (applicationDir != "") {
            if (QFile::exists(applicationDir + "/dump.txt")) {
                data = applicationDir + "/dump.txt";
            }
        }
    }
    hudson = settings.value("hudson").toString();
    bug1_link = settings.value("bug1_link").toString();
    bug1_max = settings.value("bug1_max").toInt();
    bug2_link = settings.value("bug2_link").toString();
    bug2_max = settings.value("bug2_max").toInt();
    bug3_link = settings.value("bug3_link").toString();
    bug3_max = settings.value("bug3_max").toInt();
}

MainWindow::~MainWindow() {
}

void MainWindow::closeEvent(QCloseEvent *event) {
    QSettings settings(QSettings::IniFormat, QSettings::UserScope,
            "QHudsonResults", "QHudsonResults");
    settings.setValue("data", data);
    settings.setValue("hudson", hudson);
    settings.setValue("bug1_link", bug1_link);
    settings.setValue("bug1_max", bug1_max);
    settings.setValue("bug2_link", bug2_link);
    settings.setValue("bug2_max", bug2_max);
    settings.setValue("bug3_link", bug3_link);
    settings.setValue("bug3_max", bug3_max);
    settings.sync();
    event->accept();
}

void MainWindow::about() {
    QMessageBox::about(this, QString::fromUtf8("О программе QHudsonResults"), QString::fromUtf8("
QHudsonResults

Анализ хадсоновских результатов тестов
"));
}

void MainWindow::refresh() {
    widget.table->setDisabled(true);
    widget.table->clear();
    widget.table->setRowCount(0);
    widget.table->setSortingEnabled(false);
    readFile(data);
    widget.table->setSortingEnabled(true);
    widget.table->sortByColumn(0, Qt::AscendingOrder);
    widget.table->setSelectionBehavior(QAbstractItemView::SelectRows);
    widget.table->setEditTriggers(QAbstractItemView::NoEditTriggers);
    widget.table->setEnabled(true);
}

void MainWindow::cellClicked(int row, int column) {
    QModelIndex index = widget.table->model()->index(row, 0);
    QString name = widget.table->model()->data(index).toString();
    QString text = QString("
%1
").arg(name);

    foreach (QString jobName, jobNames) {
        QString key = QString(QString("%1|%2").arg(name).arg(jobName));
        if (allTests.contains(key)) {
            text += QString("
%1
").arg(jobName);
            QList<int> testItems = allTests.value(key);
            text += "
";
            foreach (int testIndex, testItems) {
                text += "
";
                TestItem item = testsList[testIndex];
                text += "  " + item.getBuild() + "	";
                text += "  " + item.getStatusImage() + "	";
                QString message = item.getMessage();
                if (message == "null") {
                    message = QString("OK").arg(hudson).arg(jobName).arg(item.getBuild());
                } else {
                    message = QString("%4").arg(hudson).arg(jobName).arg(item.getBuild()).arg(message);
                }
//                text += "  " + item.getBugId() + "	";
                text += "  " + message + "	";
                text += "
";
            }
            text += "
";
        }
    }
    widget.text->setText(text);
}

void MainWindow::linkActivated(const QUrl& link) {
    QDesktopServices::openUrl(link);
    return;
}

void MainWindow::readFile(const QString& fileName) {
    QString testNameFilter = widget.editTestName->text();

    jobsHash.clear();
    testsHash.clear();
    testsList.clear();
    allTests.clear();
    jobNames.clear();

    int counter = 0;
    QFile file(fileName);
    file.open(QIODevice::ReadOnly);
    QTextStream  in(&file);
    while (!in.atEnd()) {
        QString line = in.readLine();

        if (line.contains(testNameFilter, Qt::CaseInsensitive)) {
            QStringList fields = line.split(",");
            QString job = fields.at(0).trimmed();
            QString build = fields.at(1).trimmed();
            QString name = fields.at(2).trimmed();
            QString status = fields.at(3).trimmed();
            QString bugId = fields.at(4).trimmed();
            QString message = fields.at(5).trimmed();

            if (name.contains(testNameFilter, Qt::CaseInsensitive)) {
                testsList.push_back(TestItem(job, build, name, status, bugId, message));
                jobsHash[job] = 0;
                testsHash[name] = 0;
                QString key = QString(QString("%1|%2").arg(name).arg(job));
                if (!allTests.contains(key)) {
                    allTests[key] = QList<int>();
                }
                allTests[key].append(counter);
                counter++;
            }
        }
    }
    file.close();

    int sizeJobs = jobsHash.size();
    widget.table->setRowCount(testsHash.size());
    widget.table->setColumnCount(sizeJobs + 1);
    widget.table->setColumnWidth(0, 400);
    widget.table->horizontalHeader()->setFixedHeight(40);
    for (int i = 1; i < sizeJobs; i++) {
        widget.table->setColumnWidth(i, 120);
    }

    QStringList headers = QStringList();
    headers.append("Имя теста");
    counter = 0;
    foreach (QString job, jobsHash.keys()) {
        headers.append(job);
        jobNames.append(job);
        jobsHash[job] = counter;
        counter++;
    }
    widget.table->setHorizontalHeaderLabels(headers);

    counter = 0;
    foreach (QString test, testsHash.keys()) {
        QTableWidgetItem *item1 = new QTableWidgetItem();
        item1->setText(test);
        widget.table->setItem(counter, 0, item1);

        int column = 0;
        foreach (QString jobName, jobNames) {
            QString key = QString(QString("%1|%2").arg(test).arg(jobName));
            QString cellValue("");
            if (allTests.contains(key)) {
                QList<int> testItems = allTests.value(key);
                foreach (int testIndex, testItems) {
                    cellValue += testsList[testIndex].getStatusChar();
                }
            }
            QTableWidgetItem *item = new QTableWidgetItem();
            item->setText(cellValue);
            widget.table->setItem(counter, column + 1, item);
            column++;
        }

        counter++;
    }
}
<commit_msg>поправил сборку на Solaris<commit_after>/*
Copyright (c) 2016, Valeriy Soldatov
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
 * Neither the name of the football.mojgorod.ru nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <QCloseEvent>
#include <QDesktopServices>
#include <QDesktopWidget>
#include <QFile>
#include <QList>
#include <QMap>
#include <QMessageBox>
#include <QSettings>
#include <QString>
#include <QTextStream>

#include "MainWindow.h"
#include "TestItem.h"

#include "ui_MainWindow.h"

using namespace std;

MainWindow::MainWindow() {

    widget.setupUi(this);

    connect(widget.actionExit, SIGNAL(triggered()),
            this, SLOT(close()));
    connect(widget.actionAbout, SIGNAL(triggered()),
            this, SLOT(about()));
    connect(widget.btnRefresh, SIGNAL(clicked()),
            this, SLOT(refresh()));
    connect(widget.table, SIGNAL(cellClicked(int, int)),
            this, SLOT(cellClicked(int, int)));
    connect(widget.text, SIGNAL(anchorClicked(const QUrl &)),
            this, SLOT(linkActivated(const QUrl &)));

    QRect rectangle = frameGeometry();
    rectangle.moveCenter(QDesktopWidget().availableGeometry().center());
    move(rectangle.topLeft());

    QSettings settings(QSettings::IniFormat, QSettings::UserScope,
            "QHudsonResults", "QHudsonResults");
    data = settings.value("data").toString();
    if (data == "") {
        QString applicationDir = QCoreApplication::applicationDirPath();
        if (applicationDir != "") {
            if (QFile::exists(applicationDir + "/dump.txt")) {
                data = applicationDir + "/dump.txt";
            }
        }
    }
    hudson = settings.value("hudson").toString();
    bug1_link = settings.value("bug1_link").toString();
    bug1_max = settings.value("bug1_max").toInt();
    bug2_link = settings.value("bug2_link").toString();
    bug2_max = settings.value("bug2_max").toInt();
    bug3_link = settings.value("bug3_link").toString();
    bug3_max = settings.value("bug3_max").toInt();
}

MainWindow::~MainWindow() {
}

void MainWindow::closeEvent(QCloseEvent *event) {
    QSettings settings(QSettings::IniFormat, QSettings::UserScope,
            "QHudsonResults", "QHudsonResults");
    settings.setValue("data", data);
    settings.setValue("hudson", hudson);
    settings.setValue("bug1_link", bug1_link);
    settings.setValue("bug1_max", bug1_max);
    settings.setValue("bug2_link", bug2_link);
    settings.setValue("bug2_max", bug2_max);
    settings.setValue("bug3_link", bug3_link);
    settings.setValue("bug3_max", bug3_max);
    settings.sync();
    event->accept();
}

void MainWindow::about() {
    QMessageBox::about(this, QString::fromUtf8("О программе QHudsonResults"), QString::fromUtf8("
QHudsonResults

Анализ хадсоновских результатов тестов
"));
}

void MainWindow::refresh() {
    widget.table->setDisabled(true);
    widget.table->clear();
    widget.table->setRowCount(0);
    widget.table->setSortingEnabled(false);
    readFile(data);
    widget.table->setSortingEnabled(true);
    widget.table->sortByColumn(0, Qt::AscendingOrder);
    widget.table->setSelectionBehavior(QAbstractItemView::SelectRows);
    widget.table->setEditTriggers(QAbstractItemView::NoEditTriggers);
    widget.table->setEnabled(true);
}

void MainWindow::cellClicked(int row, int column) {
    QModelIndex index = widget.table->model()->index(row, 0);
    QString name = widget.table->model()->data(index).toString();
    QString text = QString("
%1
").arg(name);

    foreach (QString jobName, jobNames) {
        QString key = QString(QString("%1|%2").arg(name).arg(jobName));
        if (allTests.contains(key)) {
            text += QString("
%1
").arg(jobName);
            QList<int> testItems = allTests.value(key);
            text += "
";
            foreach (int testIndex, testItems) {
                text += "
";
                TestItem item = testsList[testIndex];
                text += "  " + item.getBuild() + "	";
                text += "  " + item.getStatusImage() + "	";
                QString message = item.getMessage();
                if (message == "null") {
                    message = QString("OK").arg(hudson).arg(jobName).arg(item.getBuild());
                } else {
                    message = QString("%4").arg(hudson).arg(jobName).arg(item.getBuild()).arg(message);
                }
//                text += "  " + item.getBugId() + "	";
                text += "  " + message + "	";
                text += "
";
            }
            text += "
";
        }
    }
    widget.text->setText(text);
}

void MainWindow::linkActivated(const QUrl& link) {
    QDesktopServices::openUrl(link);
    return;
}

void MainWindow::readFile(const QString& fileName) {
    QString testNameFilter = widget.editTestName->text();

    jobsHash.clear();
    testsHash.clear();
    testsList.clear();
    allTests.clear();
    jobNames.clear();

    int counter = 0;
    QFile file(fileName);
    file.open(QIODevice::ReadOnly);
    QTextStream  in(&file);
    while (!in.atEnd()) {
        QString line = in.readLine();

        if (line.contains(testNameFilter, Qt::CaseInsensitive)) {
            QStringList fields = line.split(",");
            QString job = fields.at(0).trimmed();
            QString build = fields.at(1).trimmed();
            QString name = fields.at(2).trimmed();
            QString status = fields.at(3).trimmed();
            QString bugId = fields.at(4).trimmed();
            QString message = fields.at(5).trimmed();

            if (name.contains(testNameFilter, Qt::CaseInsensitive)) {
                testsList.push_back(TestItem(job, build, name, status, bugId, message));
                jobsHash[job] = 0;
                testsHash[name] = 0;
                QString key = QString(QString("%1|%2").arg(name).arg(job));
                if (!allTests.contains(key)) {
                    allTests[key] = QList<int>();
                }
                allTests[key].append(counter);
                counter++;
            }
        }
    }
    file.close();

    int sizeJobs = jobsHash.size();
    widget.table->setRowCount(testsHash.size());
    widget.table->setColumnCount(sizeJobs + 1);
    widget.table->setColumnWidth(0, 400);
    widget.table->horizontalHeader()->setFixedHeight(40);
    for (int i = 1; i < sizeJobs; i++) {
        widget.table->setColumnWidth(i, 120);
    }

    QStringList headers = QStringList();
    headers.append(QString::fromUtf8("Имя теста"));
    counter = 0;
    foreach (QString job, jobsHash.keys()) {
        headers.append(job);
        jobNames.append(job);
        jobsHash[job] = counter;
        counter++;
    }
    widget.table->setHorizontalHeaderLabels(headers);

    counter = 0;
    foreach (QString test, testsHash.keys()) {
        QTableWidgetItem *item1 = new QTableWidgetItem();
        item1->setText(test);
        widget.table->setItem(counter, 0, item1);

        int column = 0;
        foreach (QString jobName, jobNames) {
            QString key = QString(QString("%1|%2").arg(test).arg(jobName));
            QString cellValue("");
            if (allTests.contains(key)) {
                QList<int> testItems = allTests.value(key);
                foreach (int testIndex, testItems) {
                    cellValue += testsList[testIndex].getStatusChar();
                }
            }
            QTableWidgetItem *item = new QTableWidgetItem();
            item->setText(cellValue);
            widget.table->setItem(counter, column + 1, item);
            column++;
        }

        counter++;
    }
}
<|endoftext|>


<commit_before>
#include "Globals.h"  // NOTE: MSVC stupidness requires this to be the same across all modules

#include "MobSpawner.h"
#include "Mobs/IncludeAllMonsters.h"

cMobSpawner::cMobSpawner(cMonster::eFamily a_MonsterFamily, const std::set<eMonsterType>& a_AllowedTypes) :
	m_MonsterFamily(a_MonsterFamily),
	m_NewPack(true),
	m_MobType(mtInvalidType)
{
	for (std::set<eMonsterType>::const_iterator itr = a_AllowedTypes.begin(); itr != a_AllowedTypes.end(); ++itr)
	{
		if (cMonster::FamilyFromType(*itr) == a_MonsterFamily)
		{
			m_AllowedTypes.insert(*itr);
		}
	}
}

bool cMobSpawner::CheckPackCenter(BLOCKTYPE a_BlockType)
{
	// Packs of non-water mobs can only be centered on an air block
	// Packs of water mobs can only be centered on a water block
	if (m_MonsterFamily == cMonster::mfWater)
	{
		return IsBlockWater(a_BlockType);
	}
	else
	{
		return a_BlockType == E_BLOCK_AIR;
	}
}

void cMobSpawner::addIfAllowed(eMonsterType toAdd, std::set<eMonsterType>& toAddIn)
{
	std::set<eMonsterType>::iterator itr = m_AllowedTypes.find(toAdd);
	if (itr != m_AllowedTypes.end())
	{
		toAddIn.insert(toAdd);
	}
}

eMonsterType cMobSpawner::ChooseMobType(EMCSBiome a_Biome)
{
	std::set<eMonsterType> allowedMobs;

	if ((a_Biome == biMushroomIsland) || (a_Biome == biMushroomShore))
	{
		addIfAllowed(mtMooshroom, allowedMobs);
	}
	else if (a_Biome == biNether)
	{
		addIfAllowed(mtGhast, allowedMobs);
		addIfAllowed(mtZombiePigman, allowedMobs);
		addIfAllowed(mtMagmaCube, allowedMobs);
	}
	else if (a_Biome == biEnd)
	{
		addIfAllowed(mtEnderman, allowedMobs);
	}
	else
	{
		addIfAllowed(mtBat, allowedMobs);
		addIfAllowed(mtSpider, allowedMobs);
		addIfAllowed(mtZombie, allowedMobs);
		addIfAllowed(mtSkeleton, allowedMobs);
		addIfAllowed(mtCreeper, allowedMobs);
		addIfAllowed(mtSquid, allowedMobs);
		addIfAllowed(mtGuardian, allowedMobs);

		if ((a_Biome != biDesert) && (a_Biome != biBeach) && (a_Biome != biOcean))
		{
			addIfAllowed(mtSheep, allowedMobs);
			addIfAllowed(mtPig, allowedMobs);
			addIfAllowed(mtCow, allowedMobs);
			addIfAllowed(mtChicken, allowedMobs);
			addIfAllowed(mtEnderman, allowedMobs);
			addIfAllowed(mtRabbit, allowedMobs);
			addIfAllowed(mtSlime, allowedMobs);  // MG TODO : much more complicated rule

			if ((a_Biome == biForest) || (a_Biome == biForestHills) || (a_Biome == biTaiga) || (a_Biome == biTaigaHills))
			{
				addIfAllowed(mtWolf, allowedMobs);
			}
			else if ((a_Biome == biJungle) || (a_Biome == biJungleHills))
			{
				addIfAllowed(mtOcelot, allowedMobs);
			}
		}
	}

	size_t allowedMobsSize = allowedMobs.size();
	if (allowedMobsSize > 0)
	{
		std::set<eMonsterType>::iterator itr = allowedMobs.begin();
		int iRandom = m_Random.NextInt((int)allowedMobsSize, a_Biome);

		for (int i = 0; i < iRandom; i++)
		{
			++itr;
		}

		return *itr;
	}
	return mtInvalidType;
}

bool cMobSpawner::CanSpawnHere(cChunk * a_Chunk, int a_RelX, int a_RelY, int a_RelZ, eMonsterType a_MobType, EMCSBiome a_Biome)
{
	cFastRandom Random;
	BLOCKTYPE TargetBlock = E_BLOCK_AIR;
	if (a_Chunk->UnboundedRelGetBlockType(a_RelX, a_RelY, a_RelZ, TargetBlock))
	{
		if ((a_RelY + 1 > cChunkDef::Height) || (a_RelY - 1 < 0))
		{
			return false;
		}

		NIBBLETYPE BlockLight = a_Chunk->GetBlockLight(a_RelX, a_RelY, a_RelZ);
		NIBBLETYPE SkyLight = a_Chunk->GetSkyLight(a_RelX, a_RelY, a_RelZ);
		BLOCKTYPE BlockAbove = a_Chunk->GetBlock(a_RelX, a_RelY + 1, a_RelZ);
		BLOCKTYPE BlockBelow = a_Chunk->GetBlock(a_RelX, a_RelY - 1, a_RelZ);

		SkyLight = a_Chunk->GetTimeAlteredLight(SkyLight);

		switch (a_MobType)
		{
			case mtGuardian:
			{
				return IsBlockWater(TargetBlock) && (a_RelY >= 45) && (a_RelY <= 62);
			}

			case mtSquid:
			{
				return IsBlockWater(TargetBlock) && (a_RelY >= 45) && (a_RelY <= 62);
			}

			case mtBat:
			{
				return (a_RelY <= 63) && (BlockLight <= 4) && (SkyLight <= 4) && (TargetBlock == E_BLOCK_AIR) && !cBlockInfo::IsTransparent(BlockAbove);
			}

			case mtChicken:
			case mtCow:
			case mtPig:
			case mtHorse:
			case mtRabbit:
			case mtSheep:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(!cBlockInfo::IsTransparent(BlockBelow)) &&
					(BlockBelow == E_BLOCK_GRASS) &&
					(SkyLight >= 9)
				);
			}

			case mtOcelot:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(
						(BlockBelow == E_BLOCK_GRASS) || (BlockBelow == E_BLOCK_LEAVES) || (BlockBelow == E_BLOCK_NEW_LEAVES)
					) &&
					(a_RelY >= 62) &&
					(Random.NextInt(3, a_Biome) != 0)
				);
			}

			case mtEnderman:
			{
				if (a_RelY < 250)
				{
					BLOCKTYPE BlockTop = a_Chunk->GetBlock(a_RelX, a_RelY + 2, a_RelZ);
					if (BlockTop == E_BLOCK_AIR)
					{
						BlockTop = a_Chunk->GetBlock(a_RelX, a_RelY + 3, a_RelZ);
						return (
							(TargetBlock == E_BLOCK_AIR) &&
							(BlockAbove == E_BLOCK_AIR) &&
							(BlockTop == E_BLOCK_AIR) &&
							(!cBlockInfo::IsTransparent(BlockBelow)) &&
							(SkyLight <= 7) &&
							(BlockLight <= 7)
						);
					}
				}
				break;
			}

			case mtSpider:
			{
				bool CanSpawn = true;
				bool HasFloor = false;
				for (int x = 0; x < 2; ++x)
				{
					for (int z = 0; z < 2; ++z)
					{
						CanSpawn = a_Chunk->UnboundedRelGetBlockType(a_RelX + x, a_RelY, a_RelZ + z, TargetBlock);
						CanSpawn = CanSpawn && (TargetBlock == E_BLOCK_AIR);
						if (!CanSpawn)
						{
							return false;
						}
						HasFloor = (
							HasFloor ||
							(
								a_Chunk->UnboundedRelGetBlockType(a_RelX + x, a_RelY - 1, a_RelZ + z, TargetBlock) &&
								!cBlockInfo::IsTransparent(TargetBlock)
							)
						);
					}
				}
				return CanSpawn && HasFloor && (SkyLight <= 7) && (BlockLight <= 7);
			}

			case mtCreeper:
			case mtSkeleton:
			case mtZombie:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(!cBlockInfo::IsTransparent(BlockBelow)) &&
					(SkyLight <= 7) &&
					(BlockLight <= 7) &&
					(Random.NextInt(2, a_Biome) == 0)
				);
			}

			case mtMagmaCube:
			case mtSlime:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(!cBlockInfo::IsTransparent(BlockBelow)) &&
					(
						(a_RelY <= 40) || (a_Biome == biSwampland)
					)
				);
			}

			case mtGhast:
			case mtZombiePigman:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(!cBlockInfo::IsTransparent(BlockBelow)) &&
					(Random.NextInt(20, a_Biome) == 0)
				);
			}

			case mtWolf:
			{
				return (
					(TargetBlock == E_BLOCK_GRASS) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(
						(a_Biome == biTaiga) ||
						(a_Biome == biTaigaHills) ||
						(a_Biome == biForest) ||
						(a_Biome == biForestHills) ||
						(a_Biome == biColdTaiga) ||
						(a_Biome == biColdTaigaHills) ||
						(a_Biome == biTaigaM) ||
						(a_Biome == biMegaTaiga) ||
						(a_Biome == biMegaTaigaHills)
					)
				);
			}

			case mtMooshroom:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(BlockBelow == E_BLOCK_MYCELIUM) &&
					(
						(a_Biome == biMushroomShore) ||
						(a_Biome == biMushroomIsland)
					)
				);
			}

			default:
			{
				LOGD("MG TODO: Write spawning rule for mob type %d", a_MobType);
				return false;
			}
		}
	}
	return false;
}

cMonster* cMobSpawner::TryToSpawnHere(cChunk * a_Chunk, int a_RelX, int a_RelY, int a_RelZ, EMCSBiome a_Biome, int& a_MaxPackSize)
{
	cMonster* toReturn = nullptr;
	if (m_NewPack)
	{
		m_MobType = ChooseMobType(a_Biome);
		if (m_MobType == mtInvalidType)
		{
			return toReturn;
		}
		if (m_MobType == mtWolf)
		{
			a_MaxPackSize = 8;
		}
		else if (m_MobType == mtGhast)
		{
			a_MaxPackSize = 1;
		}
		m_NewPack = false;
	}

	// Make sure we are looking at the right chunk to spawn in
	a_Chunk = a_Chunk->GetRelNeighborChunkAdjustCoords(a_RelX, a_RelZ);

	if ((m_AllowedTypes.find(m_MobType) != m_AllowedTypes.end()) && CanSpawnHere(a_Chunk, a_RelX, a_RelY, a_RelZ, m_MobType, a_Biome))
	{
		cMonster * newMob = cMonster::NewMonsterFromType(m_MobType);
		if (newMob)
		{
			m_Spawned.insert(newMob);
		}
		toReturn = newMob;
	}
	return toReturn;
}

void cMobSpawner::NewPack()
{
	m_NewPack = true;
}

cMobSpawner::tSpawnedContainer & cMobSpawner::getSpawned(void)
{
	return m_Spawned;
}

bool cMobSpawner::CanSpawnAnything(void)
{
	return !m_AllowedTypes.empty();
}

<commit_msg>Fixed monster spawn randomness.<commit_after>
#include "Globals.h"  // NOTE: MSVC stupidness requires this to be the same across all modules

#include "MobSpawner.h"
#include "Mobs/IncludeAllMonsters.h"

cMobSpawner::cMobSpawner(cMonster::eFamily a_MonsterFamily, const std::set<eMonsterType>& a_AllowedTypes) :
	m_MonsterFamily(a_MonsterFamily),
	m_NewPack(true),
	m_MobType(mtInvalidType)
{
	for (std::set<eMonsterType>::const_iterator itr = a_AllowedTypes.begin(); itr != a_AllowedTypes.end(); ++itr)
	{
		if (cMonster::FamilyFromType(*itr) == a_MonsterFamily)
		{
			m_AllowedTypes.insert(*itr);
		}
	}
}

bool cMobSpawner::CheckPackCenter(BLOCKTYPE a_BlockType)
{
	// Packs of non-water mobs can only be centered on an air block
	// Packs of water mobs can only be centered on a water block
	if (m_MonsterFamily == cMonster::mfWater)
	{
		return IsBlockWater(a_BlockType);
	}
	else
	{
		return a_BlockType == E_BLOCK_AIR;
	}
}

void cMobSpawner::addIfAllowed(eMonsterType toAdd, std::set<eMonsterType>& toAddIn)
{
	std::set<eMonsterType>::iterator itr = m_AllowedTypes.find(toAdd);
	if (itr != m_AllowedTypes.end())
	{
		toAddIn.insert(toAdd);
	}
}

eMonsterType cMobSpawner::ChooseMobType(EMCSBiome a_Biome)
{
	std::set<eMonsterType> allowedMobs;

	if ((a_Biome == biMushroomIsland) || (a_Biome == biMushroomShore))
	{
		addIfAllowed(mtMooshroom, allowedMobs);
	}
	else if (a_Biome == biNether)
	{
		addIfAllowed(mtGhast, allowedMobs);
		addIfAllowed(mtZombiePigman, allowedMobs);
		addIfAllowed(mtMagmaCube, allowedMobs);
	}
	else if (a_Biome == biEnd)
	{
		addIfAllowed(mtEnderman, allowedMobs);
	}
	else
	{
		addIfAllowed(mtBat, allowedMobs);
		addIfAllowed(mtSpider, allowedMobs);
		addIfAllowed(mtZombie, allowedMobs);
		addIfAllowed(mtSkeleton, allowedMobs);
		addIfAllowed(mtCreeper, allowedMobs);
		addIfAllowed(mtSquid, allowedMobs);
		addIfAllowed(mtGuardian, allowedMobs);

		if ((a_Biome != biDesert) && (a_Biome != biBeach) && (a_Biome != biOcean))
		{
			addIfAllowed(mtSheep, allowedMobs);
			addIfAllowed(mtPig, allowedMobs);
			addIfAllowed(mtCow, allowedMobs);
			addIfAllowed(mtChicken, allowedMobs);
			addIfAllowed(mtEnderman, allowedMobs);
			addIfAllowed(mtRabbit, allowedMobs);
			addIfAllowed(mtSlime, allowedMobs);  // MG TODO : much more complicated rule

			if ((a_Biome == biForest) || (a_Biome == biForestHills) || (a_Biome == biTaiga) || (a_Biome == biTaigaHills))
			{
				addIfAllowed(mtWolf, allowedMobs);
			}
			else if ((a_Biome == biJungle) || (a_Biome == biJungleHills))
			{
				addIfAllowed(mtOcelot, allowedMobs);
			}
		}
	}

	size_t allowedMobsSize = allowedMobs.size();
	if (allowedMobsSize > 0)
	{
		std::set<eMonsterType>::iterator itr = allowedMobs.begin();
		static int Counter = 0;
		int iRandom = m_Random.NextInt((int)allowedMobsSize, Counter++);

		for (int i = 0; i < iRandom; i++)
		{
			++itr;
		}

		return *itr;
	}
	return mtInvalidType;
}

bool cMobSpawner::CanSpawnHere(cChunk * a_Chunk, int a_RelX, int a_RelY, int a_RelZ, eMonsterType a_MobType, EMCSBiome a_Biome)
{
	cFastRandom Random;
	BLOCKTYPE TargetBlock = E_BLOCK_AIR;
	if (a_Chunk->UnboundedRelGetBlockType(a_RelX, a_RelY, a_RelZ, TargetBlock))
	{
		if ((a_RelY + 1 > cChunkDef::Height) || (a_RelY - 1 < 0))
		{
			return false;
		}

		NIBBLETYPE BlockLight = a_Chunk->GetBlockLight(a_RelX, a_RelY, a_RelZ);
		NIBBLETYPE SkyLight = a_Chunk->GetSkyLight(a_RelX, a_RelY, a_RelZ);
		BLOCKTYPE BlockAbove = a_Chunk->GetBlock(a_RelX, a_RelY + 1, a_RelZ);
		BLOCKTYPE BlockBelow = a_Chunk->GetBlock(a_RelX, a_RelY - 1, a_RelZ);

		SkyLight = a_Chunk->GetTimeAlteredLight(SkyLight);

		switch (a_MobType)
		{
			case mtGuardian:
			{
				return IsBlockWater(TargetBlock) && (a_RelY >= 45) && (a_RelY <= 62);
			}

			case mtSquid:
			{
				return IsBlockWater(TargetBlock) && (a_RelY >= 45) && (a_RelY <= 62);
			}

			case mtBat:
			{
				return (a_RelY <= 63) && (BlockLight <= 4) && (SkyLight <= 4) && (TargetBlock == E_BLOCK_AIR) && !cBlockInfo::IsTransparent(BlockAbove);
			}

			case mtChicken:
			case mtCow:
			case mtPig:
			case mtHorse:
			case mtRabbit:
			case mtSheep:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(!cBlockInfo::IsTransparent(BlockBelow)) &&
					(BlockBelow == E_BLOCK_GRASS) &&
					(SkyLight >= 9)
				);
			}

			case mtOcelot:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(
						(BlockBelow == E_BLOCK_GRASS) || (BlockBelow == E_BLOCK_LEAVES) || (BlockBelow == E_BLOCK_NEW_LEAVES)
					) &&
					(a_RelY >= 62) &&
					(Random.NextInt(3, a_Biome) != 0)
				);
			}

			case mtEnderman:
			{
				if (a_RelY < 250)
				{
					BLOCKTYPE BlockTop = a_Chunk->GetBlock(a_RelX, a_RelY + 2, a_RelZ);
					if (BlockTop == E_BLOCK_AIR)
					{
						BlockTop = a_Chunk->GetBlock(a_RelX, a_RelY + 3, a_RelZ);
						return (
							(TargetBlock == E_BLOCK_AIR) &&
							(BlockAbove == E_BLOCK_AIR) &&
							(BlockTop == E_BLOCK_AIR) &&
							(!cBlockInfo::IsTransparent(BlockBelow)) &&
							(SkyLight <= 7) &&
							(BlockLight <= 7)
						);
					}
				}
				break;
			}

			case mtSpider:
			{
				bool CanSpawn = true;
				bool HasFloor = false;
				for (int x = 0; x < 2; ++x)
				{
					for (int z = 0; z < 2; ++z)
					{
						CanSpawn = a_Chunk->UnboundedRelGetBlockType(a_RelX + x, a_RelY, a_RelZ + z, TargetBlock);
						CanSpawn = CanSpawn && (TargetBlock == E_BLOCK_AIR);
						if (!CanSpawn)
						{
							return false;
						}
						HasFloor = (
							HasFloor ||
							(
								a_Chunk->UnboundedRelGetBlockType(a_RelX + x, a_RelY - 1, a_RelZ + z, TargetBlock) &&
								!cBlockInfo::IsTransparent(TargetBlock)
							)
						);
					}
				}
				return CanSpawn && HasFloor && (SkyLight <= 7) && (BlockLight <= 7);
			}

			case mtCreeper:
			case mtSkeleton:
			case mtZombie:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(!cBlockInfo::IsTransparent(BlockBelow)) &&
					(SkyLight <= 7) &&
					(BlockLight <= 7) &&
					(Random.NextInt(2, a_Biome) == 0)
				);
			}

			case mtMagmaCube:
			case mtSlime:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(!cBlockInfo::IsTransparent(BlockBelow)) &&
					(
						(a_RelY <= 40) || (a_Biome == biSwampland)
					)
				);
			}

			case mtGhast:
			case mtZombiePigman:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(!cBlockInfo::IsTransparent(BlockBelow)) &&
					(Random.NextInt(20, a_Biome) == 0)
				);
			}

			case mtWolf:
			{
				return (
					(TargetBlock == E_BLOCK_GRASS) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(
						(a_Biome == biTaiga) ||
						(a_Biome == biTaigaHills) ||
						(a_Biome == biForest) ||
						(a_Biome == biForestHills) ||
						(a_Biome == biColdTaiga) ||
						(a_Biome == biColdTaigaHills) ||
						(a_Biome == biTaigaM) ||
						(a_Biome == biMegaTaiga) ||
						(a_Biome == biMegaTaigaHills)
					)
				);
			}

			case mtMooshroom:
			{
				return (
					(TargetBlock == E_BLOCK_AIR) &&
					(BlockAbove == E_BLOCK_AIR) &&
					(BlockBelow == E_BLOCK_MYCELIUM) &&
					(
						(a_Biome == biMushroomShore) ||
						(a_Biome == biMushroomIsland)
					)
				);
			}

			default:
			{
				LOGD("MG TODO: Write spawning rule for mob type %d", a_MobType);
				return false;
			}
		}
	}
	return false;
}

cMonster* cMobSpawner::TryToSpawnHere(cChunk * a_Chunk, int a_RelX, int a_RelY, int a_RelZ, EMCSBiome a_Biome, int& a_MaxPackSize)
{
	cMonster* toReturn = nullptr;
	if (m_NewPack)
	{
		m_MobType = ChooseMobType(a_Biome);
		if (m_MobType == mtInvalidType)
		{
			return toReturn;
		}
		if (m_MobType == mtWolf)
		{
			a_MaxPackSize = 8;
		}
		else if (m_MobType == mtGhast)
		{
			a_MaxPackSize = 1;
		}
		m_NewPack = false;
	}

	// Make sure we are looking at the right chunk to spawn in
	a_Chunk = a_Chunk->GetRelNeighborChunkAdjustCoords(a_RelX, a_RelZ);

	if ((m_AllowedTypes.find(m_MobType) != m_AllowedTypes.end()) && CanSpawnHere(a_Chunk, a_RelX, a_RelY, a_RelZ, m_MobType, a_Biome))
	{
		cMonster * newMob = cMonster::NewMonsterFromType(m_MobType);
		if (newMob)
		{
			m_Spawned.insert(newMob);
		}
		toReturn = newMob;
	}
	return toReturn;
}

void cMobSpawner::NewPack()
{
	m_NewPack = true;
}

cMobSpawner::tSpawnedContainer & cMobSpawner::getSpawned(void)
{
	return m_Spawned;
}

bool cMobSpawner::CanSpawnAnything(void)
{
	return !m_AllowedTypes.empty();
}

<|endoftext|>


<commit_before>#include "../interface/MuonsFiller.h"

#include "DataFormats/VertexReco/interface/Vertex.h"
#include "DataFormats/MuonReco/interface/MuonSelectors.h"
#include "DataFormats/Math/interface/deltaR.h"
#include "DataFormats/PatCandidates/interface/Muon.h"
#include "DataFormats/PatCandidates/interface/TriggerObjectStandAlone.h"
#include "DataFormats/PatCandidates/interface/PackedCandidate.h"
#include "DataFormats/Common/interface/RefToPtr.h"

MuonsFiller::MuonsFiller(std::string const& _name, edm::ParameterSet const& _cfg, edm::ConsumesCollector& _coll) :
  FillerBase(_name, _cfg)
{
  getToken_(muonsToken_, _cfg, _coll, "muons");
  getToken_(verticesToken_, _cfg, _coll, "common", "vertices");

  if (useTrigger_) {
    for (unsigned iT(0); iT != panda::Muon::nTriggerObjects; ++iT) {
      std::string name(panda::Muon::TriggerObjectName[iT]); // "f<trigger filter name>"
      auto filters(getParameter_<VString>(_cfg, "triggerObjects." + name.substr(1)));
      triggerObjects_[iT].insert(filters.begin(), filters.end());
    }
  }
}

void
MuonsFiller::addOutput(TFile& _outputFile)
{
  TDirectory::TContext context(&_outputFile);
  auto* t(panda::utils::makeDocTree("MuonTriggerObject", panda::Muon::TriggerObjectName, panda::Muon::nTriggerObjects));
  t->Write();
  delete t;
}

void
MuonsFiller::branchNames(panda::utils::BranchList& _eventBranches, panda::utils::BranchList&) const
{
  _eventBranches.emplace_back("muons");

  if (isRealData_)
    _eventBranches.emplace_back("!muons.matchedGen_");
  if (!useTrigger_)
    _eventBranches.emplace_back("!muons.triggerMatch");
}

void
MuonsFiller::fill(panda::Event& _outEvent, edm::Event const& _inEvent, edm::EventSetup const& _setup)
{
  auto& inMuons(getProduct_(_inEvent, muonsToken_));
  auto& vertices(getProduct_(_inEvent, verticesToken_));

  auto& outMuons(_outEvent.muons);

  std::vector<edm::Ptr<reco::Muon>> ptrList;

  unsigned iMu(-1);
  for (auto& inMuon : inMuons) {
    ++iMu;
    auto& outMuon(outMuons.create_back());

    fillP4(outMuon, inMuon);

    outMuon.global = inMuon.isGlobalMuon();
    outMuon.tracker = inMuon.isTrackerMuon();
    outMuon.pf = inMuon.isPFMuon();

    auto&& innerTrack(inMuon.innerTrack());
    if (innerTrack.isNonnull()) {
      outMuon.validFraction = innerTrack->validFraction();
      auto&& hitPattern(innerTrack->hitPattern());
      outMuon.trkLayersWithMmt = hitPattern.trackerLayersWithMeasurement();
      outMuon.pixLayersWithMmt = hitPattern.pixelLayersWithMeasurement();
      outMuon.nValidPixel = hitPattern.numberOfValidPixelHits();
    }

    auto&& globalTrack(inMuon.globalTrack());
    if (globalTrack.isNonnull()) {
      outMuon.normChi2 = globalTrack->normalizedChi2();
      auto&& hitPattern(globalTrack->hitPattern());
      outMuon.nValidMuon = hitPattern.numberOfValidMuonHits();
    }

    outMuon.nMatched = inMuon.numberOfMatchedStations();

    auto&& combQuality(inMuon.combinedQuality());
    outMuon.chi2LocalPosition = combQuality.chi2LocalPosition;
    outMuon.trkKink = combQuality.chi2LocalPosition;

    outMuon.segmentCompatibility = muon::segmentCompatibility(inMuon);

    outMuon.charge = inMuon.charge();

    auto& pfIso(inMuon.pfIsolationR04());

    outMuon.chIso = pfIso.sumChargedHadronPt;
    outMuon.nhIso = pfIso.sumNeutralHadronEt;
    outMuon.phIso = pfIso.sumPhotonEt;
    outMuon.puIso = pfIso.sumPUPt;
    outMuon.r03Iso = inMuon.isolationR03().sumPt;

    auto* patMuon(dynamic_cast<pat::Muon const*>(&inMuon));

    if (patMuon) {
      outMuon.loose = patMuon->isLooseMuon();
      outMuon.medium = patMuon->isMediumMuon();
      // Following the "short-term instruction for Moriond 2017" given in https://twiki.cern.ch/twiki/bin/viewauth/CMS/SWGuideMuonIdRun2#MediumID2016_to_be_used_with_Run
      // Valid only for runs B-F
      outMuon.mediumBtoF = outMuon.loose && inMuon.innerTrack()->validFraction() > 0.49 &&
        ((inMuon.isGlobalMuon() &&
          inMuon.globalTrack()->normalizedChi2() < 3. &&
          inMuon.combinedQuality().chi2LocalPosition < 12. &&
          inMuon.combinedQuality().trkKink < 20. &&
          muon::segmentCompatibility(inMuon) > 0.303) ||
         muon::segmentCompatibility(inMuon) > 0.451);

      if (vertices.size() == 0) {
        outMuon.tight = false;
        outMuon.soft = false;
      }
      else {
        outMuon.tight = patMuon->isTightMuon(vertices.at(0));
        outMuon.soft = patMuon->isSoftMuon(vertices.at(0));
      }
    }
    else {
      outMuon.loose = muon::isLooseMuon(inMuon);
      outMuon.medium = muon::isMediumMuon(inMuon);
      if (vertices.size() == 0) {
        outMuon.tight = false;
        outMuon.soft = false;
      }
      else {
        outMuon.tight = muon::isTightMuon(inMuon, vertices.at(0));
        outMuon.soft = muon::isSoftMuon(inMuon, vertices.at(0));
      }
    }

    outMuon.hltsafe = outMuon.combIso() / outMuon.pt() < 0.4 && outMuon.r03Iso / outMuon.pt() < 0.4;

    auto bestTrack(inMuon.muonBestTrack());
    if (vertices.size() != 0) {
      auto& pv(vertices.at(0));
      auto pos(pv.position());
      if (patMuon)
        outMuon.dxy = patMuon->dB(); // probably gives identical value as bestTrack->dxy()
      else
        outMuon.dxy = std::abs(bestTrack->dxy(pos));

      outMuon.dz = std::abs(bestTrack->dz(pos));
    }
    else {
      if (patMuon)
        outMuon.dxy = patMuon->dB(); // probably gives identical value as bestTrack->dxy()
      else
        outMuon.dxy = std::abs(bestTrack->dxy());

      outMuon.dz = std::abs(bestTrack->dz());
    }

    outMuon.pfPt = inMuon.pfP4().pt();

    ptrList.push_back(inMuons.ptrAt(iMu));
  }

  auto originalIndices(outMuons.sort(panda::Particle::PtGreater));

  // export panda <-> reco mapping

  auto& muMuMap(objectMap_->get<reco::Muon, panda::Muon>());
  auto& pfMuMap(objectMap_->get<reco::Candidate, panda::Muon>("pf"));
  auto& vtxMuMap(objectMap_->get<reco::Vertex, panda::Muon>());
  auto& genMuMap(objectMap_->get<reco::Candidate, panda::Muon>("gen"));

  for (unsigned iP(0); iP != outMuons.size(); ++iP) {
    auto& outMuon(outMuons[iP]);
    unsigned idx(originalIndices[iP]);
    muMuMap.add(ptrList[idx], outMuon);

    auto sourcePtr(ptrList[idx]->sourceCandidatePtr(0));
    if (sourcePtr.isNonnull()) {
      pfMuMap.add(sourcePtr, outMuon);
      if (dynamic_cast<pat::PackedCandidate const*>(sourcePtr.get())) {
        auto vtxRef(static_cast<pat::PackedCandidate const&>(*sourcePtr).vertexRef());
        if (vtxRef.isNonnull())
          vtxMuMap.add(edm::refToPtr(vtxRef), outMuon);
      }
    }

    if (!isRealData_) {
      auto& inMuon(*ptrList[idx]);

      if (dynamic_cast<pat::Muon const*>(&inMuon)) {
        auto& patMuon(static_cast<pat::Muon const&>(inMuon));
        auto ref(patMuon.genParticleRef());
        if (ref.isNonnull())
          genMuMap.add(edm::refToPtr(ref), outMuon);
      }
    }
  }
}

void
MuonsFiller::setRefs(ObjectMapStore const& _objectMaps)
{
  auto& pfMuMap(objectMap_->get<reco::Candidate, panda::Muon>("pf"));
  auto& vtxMuMap(objectMap_->get<reco::Vertex, panda::Muon>());

  auto& pfMap(_objectMaps.at("pfCandidates").get<reco::Candidate, panda::PFCand>().fwdMap);
  auto& vtxMap(_objectMaps.at("vertices").get<reco::Vertex, panda::RecoVertex>().fwdMap);

  for (auto& link : pfMuMap.bwdMap) { // panda -> edm
    auto& outMuon(*link.first);
    auto& pfPtr(link.second);

    // muon sourceCandidatePtr can point to the AOD pfCandidates in some cases
    auto pfItr(pfMap.find(pfPtr));
    if (pfItr == pfMap.end())
      continue;

    outMuon.matchedPF.setRef(pfItr->second);
  }

  for (auto& link : vtxMuMap.bwdMap) { // panda -> edm
    auto& outMuon(*link.first);
    auto& vtxPtr(link.second);

    outMuon.vertex.setRef(vtxMap.at(vtxPtr));
  }

  if (!isRealData_) {
    auto& genMuMap(objectMap_->get<reco::Candidate, panda::Muon>("gen"));

    auto& genMap(_objectMaps.at("genParticles").get<reco::Candidate, panda::GenParticle>().fwdMap);

    for (auto& link : genMuMap.bwdMap) {
      auto& genPtr(link.second);
      if (genMap.find(genPtr) == genMap.end())
        continue;

      auto& outMuon(*link.first);
      outMuon.matchedGen.setRef(genMap.at(genPtr));
    }
  }

  if (useTrigger_) {
    auto& objMap(_objectMaps.at("global").get<pat::TriggerObjectStandAlone, VString>().fwdMap);

    std::vector<pat::TriggerObjectStandAlone const*> triggerObjects[panda::Muon::nTriggerObjects];

    // loop over the trigger filters we are interested in
    for (unsigned iT(0); iT != panda::Muon::nTriggerObjects; ++iT) {
      // loop over all trigger objects (and their associated filter names)
      for (auto& objAndNames : objMap) { // (TO ptr, VString)
        VString const& names(*objAndNames.second);
        // loop over the associated filter names
        for (auto& name : names) {
          if (triggerObjects_[iT].find(name) != triggerObjects_[iT].end()) {
            triggerObjects[iT].push_back(&*objAndNames.first);
            break;
          }
        }
      }
    }

    auto& muMuMap(objectMap_->get<reco::Muon, panda::Muon>().fwdMap);

    for (auto& link : muMuMap) { // edm -> panda
      auto& inMuon(*link.first);
      auto& outMuon(*link.second);

      for (unsigned iT(0); iT != panda::Muon::nTriggerObjects; ++iT) {
        for (auto* obj : triggerObjects[iT]) {
          if (reco::deltaR(inMuon, *obj) < 0.3) {
            outMuon.triggerMatch[iT] = true;
            break;
          }
        }
      }
    }
  }
}

DEFINE_TREEFILLER(MuonsFiller);
<commit_msg>bugfix<commit_after>#include "../interface/MuonsFiller.h"

#include "DataFormats/VertexReco/interface/Vertex.h"
#include "DataFormats/MuonReco/interface/MuonSelectors.h"
#include "DataFormats/Math/interface/deltaR.h"
#include "DataFormats/PatCandidates/interface/Muon.h"
#include "DataFormats/PatCandidates/interface/TriggerObjectStandAlone.h"
#include "DataFormats/PatCandidates/interface/PackedCandidate.h"
#include "DataFormats/Common/interface/RefToPtr.h"

MuonsFiller::MuonsFiller(std::string const& _name, edm::ParameterSet const& _cfg, edm::ConsumesCollector& _coll) :
  FillerBase(_name, _cfg)
{
  getToken_(muonsToken_, _cfg, _coll, "muons");
  getToken_(verticesToken_, _cfg, _coll, "common", "vertices");

  if (useTrigger_) {
    for (unsigned iT(0); iT != panda::Muon::nTriggerObjects; ++iT) {
      std::string name(panda::Muon::TriggerObjectName[iT]); // "f<trigger filter name>"
      auto filters(getParameter_<VString>(_cfg, "triggerObjects." + name.substr(1)));
      triggerObjects_[iT].insert(filters.begin(), filters.end());
    }
  }
}

void
MuonsFiller::addOutput(TFile& _outputFile)
{
  TDirectory::TContext context(&_outputFile);
  auto* t(panda::utils::makeDocTree("MuonTriggerObject", panda::Muon::TriggerObjectName, panda::Muon::nTriggerObjects));
  t->Write();
  delete t;
}

void
MuonsFiller::branchNames(panda::utils::BranchList& _eventBranches, panda::utils::BranchList&) const
{
  _eventBranches.emplace_back("muons");

  if (isRealData_)
    _eventBranches.emplace_back("!muons.matchedGen_");
  if (!useTrigger_)
    _eventBranches.emplace_back("!muons.triggerMatch");
}

void
MuonsFiller::fill(panda::Event& _outEvent, edm::Event const& _inEvent, edm::EventSetup const& _setup)
{
  auto& inMuons(getProduct_(_inEvent, muonsToken_));
  auto& vertices(getProduct_(_inEvent, verticesToken_));

  auto& outMuons(_outEvent.muons);

  std::vector<edm::Ptr<reco::Muon>> ptrList;

  unsigned iMu(-1);
  for (auto& inMuon : inMuons) {
    ++iMu;
    auto& outMuon(outMuons.create_back());

    fillP4(outMuon, inMuon);

    outMuon.global = inMuon.isGlobalMuon();
    outMuon.tracker = inMuon.isTrackerMuon();
    outMuon.pf = inMuon.isPFMuon();

    auto&& innerTrack(inMuon.innerTrack());
    if (innerTrack.isNonnull()) {
      outMuon.validFraction = innerTrack->validFraction();
      auto&& hitPattern(innerTrack->hitPattern());
      outMuon.trkLayersWithMmt = hitPattern.trackerLayersWithMeasurement();
      outMuon.pixLayersWithMmt = hitPattern.pixelLayersWithMeasurement();
      outMuon.nValidPixel = hitPattern.numberOfValidPixelHits();
    }

    auto&& globalTrack(inMuon.globalTrack());
    if (globalTrack.isNonnull()) {
      outMuon.normChi2 = globalTrack->normalizedChi2();
      auto&& hitPattern(globalTrack->hitPattern());
      outMuon.nValidMuon = hitPattern.numberOfValidMuonHits();
    }

    outMuon.nMatched = inMuon.numberOfMatchedStations();

    auto&& combQuality(inMuon.combinedQuality());
    outMuon.chi2LocalPosition = combQuality.chi2LocalPosition;
    outMuon.trkKink = combQuality.trkKink;

    outMuon.segmentCompatibility = muon::segmentCompatibility(inMuon);

    outMuon.charge = inMuon.charge();

    auto& pfIso(inMuon.pfIsolationR04());

    outMuon.chIso = pfIso.sumChargedHadronPt;
    outMuon.nhIso = pfIso.sumNeutralHadronEt;
    outMuon.phIso = pfIso.sumPhotonEt;
    outMuon.puIso = pfIso.sumPUPt;
    outMuon.r03Iso = inMuon.isolationR03().sumPt;

    auto* patMuon(dynamic_cast<pat::Muon const*>(&inMuon));

    if (patMuon) {
      outMuon.loose = patMuon->isLooseMuon();
      outMuon.medium = patMuon->isMediumMuon();
      // Following the "short-term instruction for Moriond 2017" given in https://twiki.cern.ch/twiki/bin/viewauth/CMS/SWGuideMuonIdRun2#MediumID2016_to_be_used_with_Run
      // Valid only for runs B-F
      outMuon.mediumBtoF = outMuon.loose && inMuon.innerTrack()->validFraction() > 0.49 &&
        ((inMuon.isGlobalMuon() &&
          inMuon.globalTrack()->normalizedChi2() < 3. &&
          inMuon.combinedQuality().chi2LocalPosition < 12. &&
          inMuon.combinedQuality().trkKink < 20. &&
          muon::segmentCompatibility(inMuon) > 0.303) ||
         muon::segmentCompatibility(inMuon) > 0.451);

      if (vertices.size() == 0) {
        outMuon.tight = false;
        outMuon.soft = false;
      }
      else {
        outMuon.tight = patMuon->isTightMuon(vertices.at(0));
        outMuon.soft = patMuon->isSoftMuon(vertices.at(0));
      }
    }
    else {
      outMuon.loose = muon::isLooseMuon(inMuon);
      outMuon.medium = muon::isMediumMuon(inMuon);
      if (vertices.size() == 0) {
        outMuon.tight = false;
        outMuon.soft = false;
      }
      else {
        outMuon.tight = muon::isTightMuon(inMuon, vertices.at(0));
        outMuon.soft = muon::isSoftMuon(inMuon, vertices.at(0));
      }
    }

    outMuon.hltsafe = outMuon.combIso() / outMuon.pt() < 0.4 && outMuon.r03Iso / outMuon.pt() < 0.4;

    auto bestTrack(inMuon.muonBestTrack());
    if (vertices.size() != 0) {
      auto& pv(vertices.at(0));
      auto pos(pv.position());
      if (patMuon)
        outMuon.dxy = patMuon->dB(); // probably gives identical value as bestTrack->dxy()
      else
        outMuon.dxy = std::abs(bestTrack->dxy(pos));

      outMuon.dz = std::abs(bestTrack->dz(pos));
    }
    else {
      if (patMuon)
        outMuon.dxy = patMuon->dB(); // probably gives identical value as bestTrack->dxy()
      else
        outMuon.dxy = std::abs(bestTrack->dxy());

      outMuon.dz = std::abs(bestTrack->dz());
    }

    outMuon.pfPt = inMuon.pfP4().pt();

    ptrList.push_back(inMuons.ptrAt(iMu));
  }

  auto originalIndices(outMuons.sort(panda::Particle::PtGreater));

  // export panda <-> reco mapping

  auto& muMuMap(objectMap_->get<reco::Muon, panda::Muon>());
  auto& pfMuMap(objectMap_->get<reco::Candidate, panda::Muon>("pf"));
  auto& vtxMuMap(objectMap_->get<reco::Vertex, panda::Muon>());
  auto& genMuMap(objectMap_->get<reco::Candidate, panda::Muon>("gen"));

  for (unsigned iP(0); iP != outMuons.size(); ++iP) {
    auto& outMuon(outMuons[iP]);
    unsigned idx(originalIndices[iP]);
    muMuMap.add(ptrList[idx], outMuon);

    auto sourcePtr(ptrList[idx]->sourceCandidatePtr(0));
    if (sourcePtr.isNonnull()) {
      pfMuMap.add(sourcePtr, outMuon);
      if (dynamic_cast<pat::PackedCandidate const*>(sourcePtr.get())) {
        auto vtxRef(static_cast<pat::PackedCandidate const&>(*sourcePtr).vertexRef());
        if (vtxRef.isNonnull())
          vtxMuMap.add(edm::refToPtr(vtxRef), outMuon);
      }
    }

    if (!isRealData_) {
      auto& inMuon(*ptrList[idx]);

      if (dynamic_cast<pat::Muon const*>(&inMuon)) {
        auto& patMuon(static_cast<pat::Muon const&>(inMuon));
        auto ref(patMuon.genParticleRef());
        if (ref.isNonnull())
          genMuMap.add(edm::refToPtr(ref), outMuon);
      }
    }
  }
}

void
MuonsFiller::setRefs(ObjectMapStore const& _objectMaps)
{
  auto& pfMuMap(objectMap_->get<reco::Candidate, panda::Muon>("pf"));
  auto& vtxMuMap(objectMap_->get<reco::Vertex, panda::Muon>());

  auto& pfMap(_objectMaps.at("pfCandidates").get<reco::Candidate, panda::PFCand>().fwdMap);
  auto& vtxMap(_objectMaps.at("vertices").get<reco::Vertex, panda::RecoVertex>().fwdMap);

  for (auto& link : pfMuMap.bwdMap) { // panda -> edm
    auto& outMuon(*link.first);
    auto& pfPtr(link.second);

    // muon sourceCandidatePtr can point to the AOD pfCandidates in some cases
    auto pfItr(pfMap.find(pfPtr));
    if (pfItr == pfMap.end())
      continue;

    outMuon.matchedPF.setRef(pfItr->second);
  }

  for (auto& link : vtxMuMap.bwdMap) { // panda -> edm
    auto& outMuon(*link.first);
    auto& vtxPtr(link.second);

    outMuon.vertex.setRef(vtxMap.at(vtxPtr));
  }

  if (!isRealData_) {
    auto& genMuMap(objectMap_->get<reco::Candidate, panda::Muon>("gen"));

    auto& genMap(_objectMaps.at("genParticles").get<reco::Candidate, panda::GenParticle>().fwdMap);

    for (auto& link : genMuMap.bwdMap) {
      auto& genPtr(link.second);
      if (genMap.find(genPtr) == genMap.end())
        continue;

      auto& outMuon(*link.first);
      outMuon.matchedGen.setRef(genMap.at(genPtr));
    }
  }

  if (useTrigger_) {
    auto& objMap(_objectMaps.at("global").get<pat::TriggerObjectStandAlone, VString>().fwdMap);

    std::vector<pat::TriggerObjectStandAlone const*> triggerObjects[panda::Muon::nTriggerObjects];

    // loop over the trigger filters we are interested in
    for (unsigned iT(0); iT != panda::Muon::nTriggerObjects; ++iT) {
      // loop over all trigger objects (and their associated filter names)
      for (auto& objAndNames : objMap) { // (TO ptr, VString)
        VString const& names(*objAndNames.second);
        // loop over the associated filter names
        for (auto& name : names) {
          if (triggerObjects_[iT].find(name) != triggerObjects_[iT].end()) {
            triggerObjects[iT].push_back(&*objAndNames.first);
            break;
          }
        }
      }
    }

    auto& muMuMap(objectMap_->get<reco::Muon, panda::Muon>().fwdMap);

    for (auto& link : muMuMap) { // edm -> panda
      auto& inMuon(*link.first);
      auto& outMuon(*link.second);

      for (unsigned iT(0); iT != panda::Muon::nTriggerObjects; ++iT) {
        for (auto* obj : triggerObjects[iT]) {
          if (reco::deltaR(inMuon, *obj) < 0.3) {
            outMuon.triggerMatch[iT] = true;
            break;
          }
        }
      }
    }
  }
}

DEFINE_TREEFILLER(MuonsFiller);
<|endoftext|>


<commit_before>// Copyright Hugh Perkins 2014 hughperkins at gmail
//
// This Source Code Form is subject to the terms of the Mozilla Public License,
// v. 2.0. If a copy of the MPL was not distributed with this file, You can
// obtain one at http://mozilla.org/MPL/2.0/.

#include <stdexcept>
#include <iostream>
#include <cstring>

#include "stringhelper.h"
#include "FileHelper.h"

#include "NorbLoader.h"

using namespace std;

#undef STATIC
#undef VIRTUAL
#define STATIC
#define VIRTUAL

STATIC void NorbLoader::getDimensions( std::string trainFilepath, int *p_N, int *p_numPlanes, int *p_imageSize ) {
    char*headerBytes = FileHelper::readBinaryChunk( trainFilepath, 0, 6 * 4 );
    unsigned int *headerValues = reinterpret_cast< unsigned int *>( headerBytes );

    int magic = headerValues[0];
//    std::cout << "magic: " << magic << std::endl;
    if( magic != 0x1e3d4c55 ) {
        throw std::runtime_error("magic value doesnt match expections: " + toString(magic) );
    }
    int ndim = headerValues[1];
    int N = headerValues[2];
    int numPlanes = headerValues[3];
    int imageSize = headerValues[4];
    int imageSizeRepeated = headerValues[5];
//    std::cout << "ndim " << ndim << " " << N << " " << numPlanes << " " << imageSize << " " << imageSizeRepeated << std::endl;
    checkSame( "imageSize", imageSize, imageSizeRepeated );

//    if( maxN > 0 ) {
//        N = min( maxN, N );
//    }
//    int totalLinearSize = N * numPlanes * imageSize * imageSize;
    *p_N = N;
    *p_numPlanes = numPlanes;
    *p_imageSize = imageSize;
//    *p_imagesLinearSize = totalLinearSize;
}

STATIC void NorbLoader::load( std::string trainFilepath, unsigned char *images, int *labels ) {
    load( trainFilepath, images, labels, 0, 0 );
}

STATIC void NorbLoader::load( std::string trainFilepath, unsigned char *images, int *labels, int startN, int numExamples ) {
    int N, numPlanes, imageSize;
    // I know, this could be optimized a bit, to remove the intermediate arrays...
    loadImages( images, trainFilepath, &N, &numPlanes, &imageSize, startN, numExamples );
//    int totalLinearSize = numExamples  * numPlanes * imageSize * imageSize;
//    memcpy( images, loadedImages + startN * numPlanes * imageSize * imageSize, numExamples * numPlanes * imageSize * imageSize * sizeof( unsigned char ) );
    loadLabels( labels, replace( trainFilepath, "-dat.mat","-cat.mat"), startN, numExamples );
//    memcpy( labels, loadedLabels + startN, sizeof( int ) * numExamples );
//    delete []loadedImages;
//    delete[] loadedLabels;
}

//STATIC unsigned char *NorbLoader::loadImages( std::string filepath, int *p_N, int *p_numPlanes, int *p_imageSize ) {
//    return loadImages( filepath, p_N, p_numPlanes, p_imageSize, 0, 0 );
//}

STATIC int *NorbLoader::loadLabels( std::string labelsfilepath, int numExamples ) {
    int *labels = new int[numExamples];
    loadLabels( labels, labelsfilepath, 0, numExamples );
    return labels;
}

STATIC unsigned char *NorbLoader::loadImages( std::string filepath, int *p_N, int *p_numPlanes, int *p_imageSize ) {
    return loadImages( filepath, p_N, p_numPlanes, p_imageSize, 0, 0 );
}

STATIC unsigned char *NorbLoader::loadImages( std::string filepath, int *p_N, int *p_numPlanes, int *p_imageSize, int numExamples ) {
    return loadImages( filepath, p_N, p_numPlanes, p_imageSize, 0, numExamples );
}

STATIC unsigned char *NorbLoader::loadImages( std::string filepath, int *p_N, int *p_numPlanes, int *p_imageSize, int startN, int numExamples ) {
    getDimensions( filepath, p_N, p_numPlanes, p_imageSize );
    unsigned char *images = new unsigned char[ (long)numExamples * *p_numPlanes * *p_imageSize * *p_imageSize ];
    loadImages( images, filepath, p_N, p_numPlanes, p_imageSize, startN, numExamples );
    return images;
}

// you need to allocate this yourself, before use
STATIC void NorbLoader::loadImages( unsigned char *images, std::string filepath, int *p_N, int *p_numPlanes, int *p_imageSize, int startN, int numExamples ) {
    char*headerBytes = FileHelper::readBinaryChunk( filepath, 0, 6 * 4 );
    unsigned int *headerValues = reinterpret_cast< unsigned int *>( headerBytes );

    int magic = headerValues[0];
//    std::cout << "magic: " << magic << std::endl;
    if( magic != 0x1e3d4c55 ) {
        throw std::runtime_error("magic value doesnt match expections: " + toString(magic) );
    }
    int ndim = headerValues[1];
    int N = headerValues[2];
    int numPlanes = headerValues[3];
    int imageSize = headerValues[4];
    int imageSizeRepeated = headerValues[5];
//    std::cout << "ndim " << ndim << " " << N << " " << numPlanes << " " << imageSize << " " << imageSizeRepeated << std::endl;
    checkSame( "imageSize", imageSize, imageSizeRepeated );

    if( numExamples > 0 && numExamples > ( N - startN ) ) {
        throw runtime_error("You requested " + toString( numExamples ) + " but there are only " + toString( N - startN ) + " avialalbe after start N " + toString( startN ) );
    }
    if( numExamples == 0 ) {
        numExamples = N - startN;
    }
//    if( maxN > 0 ) {
//        N = min( maxN, N );
//    }
    long fileStartPos = 6 * 4 + (long)startN * numPlanes * imageSize * imageSize;
    long fileReadLength = (long)numExamples * numPlanes * imageSize * imageSize;
    char *imagesAsCharArray = reinterpret_cast< char *>(images );
//    cout << "images, filestartpos " << fileStartPos << " readlength " << fileReadLength << endl;
    FileHelper::readBinaryChunk( imagesAsCharArray, filepath, fileStartPos, fileReadLength );
//    unsigned char *imagesDataUnsigned = reinterpret_cast< unsigned char *>(imagesDataSigned);

    *p_N = N;
    *p_numPlanes = numPlanes;
    *p_imageSize = imageSize;
//    return imagesDataUnsigned;
}
// you need to allocate this yourself, before use
STATIC void NorbLoader::loadLabels( int *labels, std::string filepath, int startN, int numExamples ) {
    char*headerBytes = FileHelper::readBinaryChunk( filepath, 0, 6 * 5 );
    unsigned int *headerValues = reinterpret_cast< unsigned int *>( headerBytes );

    int magic = headerValues[0];
//    std::cout << "magic: " << magic << std::endl;
    if( magic != 0x1e3d4c54 ) {
        throw std::runtime_error("magic value doesnt match expections: " + toString(magic) + " expected: " + toString( 0x1e3d4c54 ) );
    }
    int ndim = headerValues[1];
    int N = headerValues[2];
//    checkSame( "ndim", 1, ndim );
    if( numExamples > 0 && numExamples > ( N - startN ) ) {
        throw runtime_error("You requested " + toString( numExamples ) + " but there are only " + toString( N - startN ) + " avialalbe after start N " + toString( startN ) );
    }
    if( numExamples == 0 ) {
        numExamples = N - startN;
    }
//    N = Ntoget;

//    int totalLinearSize = N;
    char *labelsAsCharArray = reinterpret_cast< char *>( labels );
    long fileStartPos = 5 * 4 + (long)startN * 4;
    long fileReadLength = (long)numExamples * 4;
//    cout << "labels file read start " << fileStartPos << " length " << fileReadLength << endl;
    FileHelper::readBinaryChunk( labelsAsCharArray, filepath, fileStartPos, fileReadLength );
//    int *labels = reinterpret_cast< int *>(labelsAsByteArray);
//    return labels;
}
STATIC void NorbLoader::writeImages( std::string filepath, unsigned char *images, int N, int numPlanes, int imageSize ) {
    int totalLinearSize = N * numPlanes * imageSize * imageSize;

    long imagesFilesize = totalLinearSize + 6 * 4; // magic, plus num dimensions, plus 4 dimensions
    char *imagesDataSigned = new char[ imagesFilesize ];
    unsigned int *imagesDataInt = reinterpret_cast< unsigned int *>( imagesDataSigned );
    unsigned char *imagesDataUnsigned = reinterpret_cast< unsigned char *>(imagesDataSigned);
    imagesDataInt[0] = 0x1e3d4c55;
    imagesDataInt[1] = 4;
    imagesDataInt[2] = N;
    imagesDataInt[3] = numPlanes;
    imagesDataInt[4] = imageSize;
    imagesDataInt[5] = imageSize;
    memcpy( imagesDataUnsigned + 6 * sizeof(int), images, totalLinearSize * sizeof( unsigned char ) );
    FileHelper::writeBinary( filepath, imagesDataSigned, imagesFilesize );
}
STATIC void NorbLoader::writeLabels( std::string filepath, int *labels, int N ) {
    int totalLinearSize = N;

    long imagesFilesize = totalLinearSize * 4 + 5 * 4; // magic, plus num dimensions, plus 3 dimensions
    char *imagesDataSigned = new char[ imagesFilesize ];
    unsigned int *imagesDataInt = reinterpret_cast< unsigned int *>( imagesDataSigned );
    unsigned char *imagesDataUnsigned = reinterpret_cast< unsigned char *>(imagesDataSigned);
    imagesDataInt[0] = 0x1e3d4c54;
    imagesDataInt[1] = 1;
    imagesDataInt[2] = N;
    imagesDataInt[3] = 1;
    imagesDataInt[4] = 1;
    memcpy( imagesDataUnsigned + 5 * sizeof(int), labels, totalLinearSize * sizeof( int ) );
    FileHelper::writeBinary( filepath, imagesDataSigned, imagesFilesize );
}

<commit_msg>couple more warnings removed<commit_after>// Copyright Hugh Perkins 2014 hughperkins at gmail
//
// This Source Code Form is subject to the terms of the Mozilla Public License,
// v. 2.0. If a copy of the MPL was not distributed with this file, You can
// obtain one at http://mozilla.org/MPL/2.0/.

#include <stdexcept>
#include <iostream>
#include <cstring>

#include "stringhelper.h"
#include "FileHelper.h"

#include "NorbLoader.h"

using namespace std;

#undef STATIC
#undef VIRTUAL
#define STATIC
#define VIRTUAL

STATIC void NorbLoader::getDimensions( std::string trainFilepath, int *p_N, int *p_numPlanes, int *p_imageSize ) {
    char*headerBytes = FileHelper::readBinaryChunk( trainFilepath, 0, 6 * 4 );
    unsigned int *headerValues = reinterpret_cast< unsigned int *>( headerBytes );

    int magic = headerValues[0];
//    std::cout << "magic: " << magic << std::endl;
    if( magic != 0x1e3d4c55 ) {
        throw std::runtime_error("magic value doesnt match expections: " + toString(magic) );
    }
//    int ndim = headerValues[1];
    int N = headerValues[2];
    int numPlanes = headerValues[3];
    int imageSize = headerValues[4];
    int imageSizeRepeated = headerValues[5];
//    std::cout << "ndim " << ndim << " " << N << " " << numPlanes << " " << imageSize << " " << imageSizeRepeated << std::endl;
    checkSame( "imageSize", imageSize, imageSizeRepeated );

//    if( maxN > 0 ) {
//        N = min( maxN, N );
//    }
//    int totalLinearSize = N * numPlanes * imageSize * imageSize;
    *p_N = N;
    *p_numPlanes = numPlanes;
    *p_imageSize = imageSize;
//    *p_imagesLinearSize = totalLinearSize;
}

STATIC void NorbLoader::load( std::string trainFilepath, unsigned char *images, int *labels ) {
    load( trainFilepath, images, labels, 0, 0 );
}

STATIC void NorbLoader::load( std::string trainFilepath, unsigned char *images, int *labels, int startN, int numExamples ) {
    int N, numPlanes, imageSize;
    // I know, this could be optimized a bit, to remove the intermediate arrays...
    loadImages( images, trainFilepath, &N, &numPlanes, &imageSize, startN, numExamples );
//    int totalLinearSize = numExamples  * numPlanes * imageSize * imageSize;
//    memcpy( images, loadedImages + startN * numPlanes * imageSize * imageSize, numExamples * numPlanes * imageSize * imageSize * sizeof( unsigned char ) );
    loadLabels( labels, replace( trainFilepath, "-dat.mat","-cat.mat"), startN, numExamples );
//    memcpy( labels, loadedLabels + startN, sizeof( int ) * numExamples );
//    delete []loadedImages;
//    delete[] loadedLabels;
}

//STATIC unsigned char *NorbLoader::loadImages( std::string filepath, int *p_N, int *p_numPlanes, int *p_imageSize ) {
//    return loadImages( filepath, p_N, p_numPlanes, p_imageSize, 0, 0 );
//}

STATIC int *NorbLoader::loadLabels( std::string labelsfilepath, int numExamples ) {
    int *labels = new int[numExamples];
    loadLabels( labels, labelsfilepath, 0, numExamples );
    return labels;
}

STATIC unsigned char *NorbLoader::loadImages( std::string filepath, int *p_N, int *p_numPlanes, int *p_imageSize ) {
    return loadImages( filepath, p_N, p_numPlanes, p_imageSize, 0, 0 );
}

STATIC unsigned char *NorbLoader::loadImages( std::string filepath, int *p_N, int *p_numPlanes, int *p_imageSize, int numExamples ) {
    return loadImages( filepath, p_N, p_numPlanes, p_imageSize, 0, numExamples );
}

STATIC unsigned char *NorbLoader::loadImages( std::string filepath, int *p_N, int *p_numPlanes, int *p_imageSize, int startN, int numExamples ) {
    getDimensions( filepath, p_N, p_numPlanes, p_imageSize );
    unsigned char *images = new unsigned char[ (long)numExamples * *p_numPlanes * *p_imageSize * *p_imageSize ];
    loadImages( images, filepath, p_N, p_numPlanes, p_imageSize, startN, numExamples );
    return images;
}

// you need to allocate this yourself, before use
STATIC void NorbLoader::loadImages( unsigned char *images, std::string filepath, int *p_N, int *p_numPlanes, int *p_imageSize, int startN, int numExamples ) {
    char*headerBytes = FileHelper::readBinaryChunk( filepath, 0, 6 * 4 );
    unsigned int *headerValues = reinterpret_cast< unsigned int *>( headerBytes );

    int magic = headerValues[0];
//    std::cout << "magic: " << magic << std::endl;
    if( magic != 0x1e3d4c55 ) {
        throw std::runtime_error("magic value doesnt match expections: " + toString(magic) );
    }
//    int ndim = headerValues[1];
    int N = headerValues[2];
    int numPlanes = headerValues[3];
    int imageSize = headerValues[4];
    int imageSizeRepeated = headerValues[5];
//    std::cout << "ndim " << ndim << " " << N << " " << numPlanes << " " << imageSize << " " << imageSizeRepeated << std::endl;
    checkSame( "imageSize", imageSize, imageSizeRepeated );

    if( numExamples > 0 && numExamples > ( N - startN ) ) {
        throw runtime_error("You requested " + toString( numExamples ) + " but there are only " + toString( N - startN ) + " avialalbe after start N " + toString( startN ) );
    }
    if( numExamples == 0 ) {
        numExamples = N - startN;
    }
//    if( maxN > 0 ) {
//        N = min( maxN, N );
//    }
    long fileStartPos = 6 * 4 + (long)startN * numPlanes * imageSize * imageSize;
    long fileReadLength = (long)numExamples * numPlanes * imageSize * imageSize;
    char *imagesAsCharArray = reinterpret_cast< char *>(images );
//    cout << "images, filestartpos " << fileStartPos << " readlength " << fileReadLength << endl;
    FileHelper::readBinaryChunk( imagesAsCharArray, filepath, fileStartPos, fileReadLength );
//    unsigned char *imagesDataUnsigned = reinterpret_cast< unsigned char *>(imagesDataSigned);

    *p_N = N;
    *p_numPlanes = numPlanes;
    *p_imageSize = imageSize;
//    return imagesDataUnsigned;
}
// you need to allocate this yourself, before use
STATIC void NorbLoader::loadLabels( int *labels, std::string filepath, int startN, int numExamples ) {
    char*headerBytes = FileHelper::readBinaryChunk( filepath, 0, 6 * 5 );
    unsigned int *headerValues = reinterpret_cast< unsigned int *>( headerBytes );

    int magic = headerValues[0];
//    std::cout << "magic: " << magic << std::endl;
    if( magic != 0x1e3d4c54 ) {
        throw std::runtime_error("magic value doesnt match expections: " + toString(magic) + " expected: " + toString( 0x1e3d4c54 ) );
    }
//    int ndim = headerValues[1];
    int N = headerValues[2];
//    checkSame( "ndim", 1, ndim );
    if( numExamples > 0 && numExamples > ( N - startN ) ) {
        throw runtime_error("You requested " + toString( numExamples ) + " but there are only " + toString( N - startN ) + " avialalbe after start N " + toString( startN ) );
    }
    if( numExamples == 0 ) {
        numExamples = N - startN;
    }
//    N = Ntoget;

//    int totalLinearSize = N;
    char *labelsAsCharArray = reinterpret_cast< char *>( labels );
    long fileStartPos = 5 * 4 + (long)startN * 4;
    long fileReadLength = (long)numExamples * 4;
//    cout << "labels file read start " << fileStartPos << " length " << fileReadLength << endl;
    FileHelper::readBinaryChunk( labelsAsCharArray, filepath, fileStartPos, fileReadLength );
//    int *labels = reinterpret_cast< int *>(labelsAsByteArray);
//    return labels;
}
STATIC void NorbLoader::writeImages( std::string filepath, unsigned char *images, int N, int numPlanes, int imageSize ) {
    int totalLinearSize = N * numPlanes * imageSize * imageSize;

    long imagesFilesize = totalLinearSize + 6 * 4; // magic, plus num dimensions, plus 4 dimensions
    char *imagesDataSigned = new char[ imagesFilesize ];
    unsigned int *imagesDataInt = reinterpret_cast< unsigned int *>( imagesDataSigned );
    unsigned char *imagesDataUnsigned = reinterpret_cast< unsigned char *>(imagesDataSigned);
    imagesDataInt[0] = 0x1e3d4c55;
    imagesDataInt[1] = 4;
    imagesDataInt[2] = N;
    imagesDataInt[3] = numPlanes;
    imagesDataInt[4] = imageSize;
    imagesDataInt[5] = imageSize;
    memcpy( imagesDataUnsigned + 6 * sizeof(int), images, totalLinearSize * sizeof( unsigned char ) );
    FileHelper::writeBinary( filepath, imagesDataSigned, imagesFilesize );
}
STATIC void NorbLoader::writeLabels( std::string filepath, int *labels, int N ) {
    int totalLinearSize = N;

    long imagesFilesize = totalLinearSize * 4 + 5 * 4; // magic, plus num dimensions, plus 3 dimensions
    char *imagesDataSigned = new char[ imagesFilesize ];
    unsigned int *imagesDataInt = reinterpret_cast< unsigned int *>( imagesDataSigned );
    unsigned char *imagesDataUnsigned = reinterpret_cast< unsigned char *>(imagesDataSigned);
    imagesDataInt[0] = 0x1e3d4c54;
    imagesDataInt[1] = 1;
    imagesDataInt[2] = N;
    imagesDataInt[3] = 1;
    imagesDataInt[4] = 1;
    memcpy( imagesDataUnsigned + 5 * sizeof(int), labels, totalLinearSize * sizeof( int ) );
    FileHelper::writeBinary( filepath, imagesDataSigned, imagesFilesize );
}

<|endoftext|>


<commit_before>#include "StdAfx.h"
#include "CameraMetaData.h"

CameraMetaData::CameraMetaData(char *docname)
{
  ctxt = xmlNewParserCtxt();
  doc = xmlCtxtReadFile(ctxt, docname, NULL, XML_PARSE_DTDVALID);

  if (doc == NULL ) {
    ThrowCME("CameraMetaData: XML Document could not be parsed successfully. Error was: %s", ctxt->lastError.message);
  }

  if (ctxt->valid == 0) {
    if (ctxt->lastError.code ==0x5e) {
      printf("CameraMetaData: Unable to locate DTD, attempting to ignore.");
    } else {
      ThrowCME("CameraMetaData: XML file does not validate. DTD Error was: %s", ctxt->lastError.message);
    }
  }

  xmlNodePtr cur;
  cur = xmlDocGetRootElement(doc);
  if (xmlStrcmp(cur->name, (const xmlChar *) "Cameras")) {
    ThrowCME("CameraMetaData: XML document of the wrong type, root node is not cameras.");
    return;
  }

  cur = cur->xmlChildrenNode;
  while (cur != NULL) {
    if ((!xmlStrcmp(cur->name, (const xmlChar *)"Camera"))){
      Camera *camera = new Camera(doc, cur);
      string id = string(camera->make).append(camera->model).append(camera->mode);
      if (cameras.end() !=cameras.find(id))
        printf("CameraMetaData: Duplicate entry found for camera: %s %s, Skipping!\n",camera->make.c_str(), camera->model.c_str());
      else
        cameras[id] = camera;
    }
    cur = cur->next;
  }
}

CameraMetaData::~CameraMetaData(void) {
  if (doc)
    xmlFreeDoc(doc);
  doc = 0;
  if (ctxt)
    xmlFreeParserCtxt(ctxt);
  ctxt = 0;
}

void CameraMetaData::dumpXML()
{
  map<string,Camera*>::iterator i = cameras.begin();
  for ( ; i != cameras.end(); i++) {
    dumpCameraXML((*i).second);
  }
}

void CameraMetaData::dumpCameraXML( Camera* cam )
{
  cout << "<Camera make=\"" << cam->make << "\" model = \"" << cam->model << "\">" << endl;
  cout << "<CFA width=\"2\" height=\"2\">" << endl;
  cout << "<Color x=\"0\" y=\"0\">" << ColorFilterArray::colorToString(cam->cfa.getColorAt(0,0)) << "</Color>";
  cout << "<Color x=\"1\" y=\"0\">" << ColorFilterArray::colorToString(cam->cfa.getColorAt(1,0)) << "</Color>" << endl;
  cout << "<Color x=\"0\" y=\"1\">" << ColorFilterArray::colorToString(cam->cfa.getColorAt(0,1)) << "</Color>";
  cout << "<Color x=\"1\" y=\"1\">" << ColorFilterArray::colorToString(cam->cfa.getColorAt(1,1)) << "</Color>" << endl;
  cout << "</CFA>" << endl;
  cout << "<Crop x=\"" << cam->cropPos.x << "\" y=\"" << cam->cropPos.y << "\" ";
  cout << "width=\"" << cam->cropSize.x << "\" height=\"" << cam->cropSize.y << "\"/>" << endl;
  cout << "<Sensor black=\"" <<cam->black << "\" white=\"" << cam->white << "\"/>" << endl;
  if (!cam->blackAreas.empty()) {
    cout << "<BlackAreas>" << endl;
    for (guint i = 0; i < cam->blackAreas.size(); i++) {
      BlackArea b = cam->blackAreas[i];
      if (b.isVertical) {
        cout << "<Vertical x=\"" << b.offset << "\" width=\"" << b.size << "\"/>"<< endl;
      } else {
        cout << "<Horizontal y=\"" << b.offset << "\" height=\"" << b.size << "\"/>"<< endl;
      }
    }
    cout << "</BlackAreas>" << endl;
  }
  cout << "</Camera>" <<endl;
}

Camera* CameraMetaData::getCamera( string make, string model, string mode )
{
  string id = string(make).append(model).append(mode);
  if (cameras.end() ==cameras.find(id))
    return NULL;
  return cameras[id];
}<commit_msg>- Fix metadata leak.<commit_after>#include "StdAfx.h"
#include "CameraMetaData.h"

CameraMetaData::CameraMetaData(char *docname)
{
  ctxt = xmlNewParserCtxt();
  doc = xmlCtxtReadFile(ctxt, docname, NULL, XML_PARSE_DTDVALID);

  if (doc == NULL ) {
    ThrowCME("CameraMetaData: XML Document could not be parsed successfully. Error was: %s", ctxt->lastError.message);
  }

  if (ctxt->valid == 0) {
    if (ctxt->lastError.code ==0x5e) {
      printf("CameraMetaData: Unable to locate DTD, attempting to ignore.");
    } else {
      ThrowCME("CameraMetaData: XML file does not validate. DTD Error was: %s", ctxt->lastError.message);
    }
  }

  xmlNodePtr cur;
  cur = xmlDocGetRootElement(doc);
  if (xmlStrcmp(cur->name, (const xmlChar *) "Cameras")) {
    ThrowCME("CameraMetaData: XML document of the wrong type, root node is not cameras.");
    return;
  }

  cur = cur->xmlChildrenNode;
  while (cur != NULL) {
    if ((!xmlStrcmp(cur->name, (const xmlChar *)"Camera"))){
      Camera *camera = new Camera(doc, cur);
      string id = string(camera->make).append(camera->model).append(camera->mode);
      if (cameras.end() !=cameras.find(id))
        printf("CameraMetaData: Duplicate entry found for camera: %s %s, Skipping!\n",camera->make.c_str(), camera->model.c_str());
      else
        cameras[id] = camera;
    }
    cur = cur->next;
  }
}

CameraMetaData::~CameraMetaData(void) {
  map<string,Camera*>::iterator i = cameras.begin();
  for ( ; i != cameras.end(); i++) {
    delete((*i).second);
  }
  if (doc)
    xmlFreeDoc(doc);
  doc = 0;
  if (ctxt)
    xmlFreeParserCtxt(ctxt);
  ctxt = 0;
}

void CameraMetaData::dumpXML()
{
  map<string,Camera*>::iterator i = cameras.begin();
  for ( ; i != cameras.end(); i++) {
    dumpCameraXML((*i).second);
  }
}

void CameraMetaData::dumpCameraXML( Camera* cam )
{
  cout << "<Camera make=\"" << cam->make << "\" model = \"" << cam->model << "\">" << endl;
  cout << "<CFA width=\"2\" height=\"2\">" << endl;
  cout << "<Color x=\"0\" y=\"0\">" << ColorFilterArray::colorToString(cam->cfa.getColorAt(0,0)) << "</Color>";
  cout << "<Color x=\"1\" y=\"0\">" << ColorFilterArray::colorToString(cam->cfa.getColorAt(1,0)) << "</Color>" << endl;
  cout << "<Color x=\"0\" y=\"1\">" << ColorFilterArray::colorToString(cam->cfa.getColorAt(0,1)) << "</Color>";
  cout << "<Color x=\"1\" y=\"1\">" << ColorFilterArray::colorToString(cam->cfa.getColorAt(1,1)) << "</Color>" << endl;
  cout << "</CFA>" << endl;
  cout << "<Crop x=\"" << cam->cropPos.x << "\" y=\"" << cam->cropPos.y << "\" ";
  cout << "width=\"" << cam->cropSize.x << "\" height=\"" << cam->cropSize.y << "\"/>" << endl;
  cout << "<Sensor black=\"" <<cam->black << "\" white=\"" << cam->white << "\"/>" << endl;
  if (!cam->blackAreas.empty()) {
    cout << "<BlackAreas>" << endl;
    for (guint i = 0; i < cam->blackAreas.size(); i++) {
      BlackArea b = cam->blackAreas[i];
      if (b.isVertical) {
        cout << "<Vertical x=\"" << b.offset << "\" width=\"" << b.size << "\"/>"<< endl;
      } else {
        cout << "<Horizontal y=\"" << b.offset << "\" height=\"" << b.size << "\"/>"<< endl;
      }
    }
    cout << "</BlackAreas>" << endl;
  }
  cout << "</Camera>" <<endl;
}

Camera* CameraMetaData::getCamera( string make, string model, string mode )
{
  string id = string(make).append(model).append(mode);
  if (cameras.end() ==cameras.find(id))
    return NULL;
  return cameras[id];
}<|endoftext|>


<commit_before>/*
 *      Author: Christian Rauch
 *
 * Copyright (c) 2018, University Of Edinburgh
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *  * Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of  nor the names of its contributors may be used to
 *    endorse or promote products derived from this software without specific
 *    prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 */

#include "task_map/Point2Line.hpp"

REGISTER_TASKMAP_TYPE("Point2Line", exotica::Point2Line);

namespace exotica
{
Eigen::Vector3d Point2Line::direction(const Eigen::Vector3d &point)
{
    // http://mathworld.wolfram.com/Point-LineDistance3-Dimensional.html
    // let:
    //      s: start of line
    //      e: end of line
    //      p: point
    // then the point on vector v = e-s that is closest to p is vp = s + t*(e-s) with
    // t = -((s-p)*(e-s)) / (|e-s|^2)    (* denotes the dot product)
    if (debug_) HIGHLIGHT_NAMED("P2L", "\e[4m" << link_name << "\e[0m");
    if (debug_) HIGHLIGHT_NAMED("P2L", "p " << point.transpose());
    if (debug_) HIGHLIGHT_NAMED("P2L", "ls " << line_start.transpose());
    if (debug_) HIGHLIGHT_NAMED("P2L", "le " << line_end.transpose());
    double t = -(line_start - point).dot(line) / line.squaredNorm();
    std::stringstream ss;
    ss << "t " << t;
    if (!infinite)
    {
        // clip to to range [0,1]
        t = std::min(std::max(0.0, t), 1.0);
        ss << ", clipped |t| " << t;
    }
    if (debug_) HIGHLIGHT_NAMED("P2L", ss.str());

    // vector from point 'p' to point 'vp' on line
    // vp = line_start + t * (line_end-line_start)
    const Eigen::Vector3d dv = line_start + (t * line) - point;
    if (debug_) HIGHLIGHT_NAMED("P2L", "vp " << (line_start + (t * line)).transpose());
    if (debug_) HIGHLIGHT_NAMED("P2L", "dv " << dv.transpose());
    return dv;
}

void Point2Line::update(Eigen::VectorXdRefConst x, Eigen::VectorXdRef phi)
{
    if (phi.rows() != Kinematics[0].Phi.rows() * 3) throw_named("Wrong size of phi!");

    for (int i = 0; i < Kinematics[0].Phi.rows(); i++)
    {
        const Eigen::Vector3d p = line_start + Eigen::Map<const Eigen::Vector3d>(Kinematics[0].Phi(i).p.data);
        phi.segment<3>(i * 3) = -direction(p);
    }
}

void Point2Line::update(Eigen::VectorXdRefConst x, Eigen::VectorXdRef phi, Eigen::MatrixXdRef J)
{
    if (phi.rows() != Kinematics[0].Phi.rows() * 3) throw_named("Wrong size of phi!");
    if (J.rows() != Kinematics[0].J.rows() * 3 || J.cols() != Kinematics[0].J(0).data.cols()) throw_named("Wrong size of J! " << Kinematics[0].J(0).data.cols());

    for (int i = 0; i < Kinematics[0].Phi.rows(); i++)
    {
        // point in base frame
        const Eigen::Vector3d p = line_start + Eigen::Map<const Eigen::Vector3d>(Kinematics[0].Phi(i).p.data);
        // direction from point to line
        const Eigen::Vector3d dv = direction(p);
        phi.segment<3>(i * 3) = dv;
        for (int j = 0; j < J.cols(); j++)
        {
            J.middleRows<3>(i * 3).col(j) = Kinematics[0].J[i].data.topRows<3>().col(j).dot(line / line.squaredNorm()) * line - Kinematics[0].J[i].data.topRows<3>().col(j);
        }

        // visualisation of point, line and their distance
        if (visualise && Server::isRos())
        {
            const ros::Time t = ros::Time::now();
            const std::string common_frame = "exotica/" + base_name;
            visualization_msgs::MarkerArray ma;
            {
                // line in base frame
                visualization_msgs::Marker mc;
                mc.header.stamp = t;
                mc.frame_locked = true;
                mc.header.frame_id = common_frame;
                mc.ns = "cam/line/" + object_name_;
                mc.type = visualization_msgs::Marker::ARROW;
                mc.scale.x = 0.01;
                mc.scale.y = 0.01;
                mc.scale.z = 0.01;
                // line start
                geometry_msgs::Point pp;
                pp.x = line_start.x();
                pp.y = line_start.y();
                pp.z = line_start.z();
                mc.points.push_back(pp);
                // line end
                const Eigen::Vector3d pe = p + dv;
                pp.x = pe.x();
                pp.y = pe.y();
                pp.z = pe.z();
                mc.points.push_back(pp);
                mc.color.r = 1;
                mc.color.g = 1;
                mc.color.b = 0;
                mc.color.a = 1;
                ma.markers.push_back(mc);
            }
            {
                // point in link frame
                visualization_msgs::Marker ml;
                ml.header.stamp = t;
                ml.frame_locked = true;
                ml.header.frame_id = common_frame;
                ml.ns = "lnk/point/" + object_name_;
                ml.type = visualization_msgs::Marker::SPHERE;
                ml.scale.x = 0.03;
                ml.scale.y = 0.03;
                ml.scale.z = 0.03;
                ml.color.r = 1;
                ml.color.g = 0;
                ml.color.b = 0;
                ml.color.a = 1;
                ml.pose.position.x = p.x();
                ml.pose.position.y = p.y();
                ml.pose.position.z = p.z();
                ma.markers.push_back(ml);
            }
            {
                // draw 'dv' starting at 'p' in base frame
                visualization_msgs::Marker mdv;
                mdv.header.stamp = t;
                mdv.frame_locked = true;
                mdv.header.frame_id = common_frame;
                mdv.ns = "dv/" + object_name_;
                mdv.type = visualization_msgs::Marker::ARROW;
                mdv.scale.x = 0.001;
                mdv.scale.y = 0.01;
                mdv.scale.z = 0.01;
                mdv.pose.position.x = p.x();
                mdv.pose.position.y = p.y();
                mdv.pose.position.z = p.z();
                mdv.points.push_back(geometry_msgs::Point());
                geometry_msgs::Point pdv;
                pdv.x = dv.x();
                pdv.y = dv.y();
                pdv.z = dv.z();
                mdv.points.push_back(pdv);
                mdv.color.r = 0;
                mdv.color.g = 1;
                mdv.color.b = 0;
                mdv.color.a = 0.5;
                ma.markers.push_back(mdv);
            }
            pub_marker.publish(ma);
            {
                visualization_msgs::MarkerArray ma;
                visualization_msgs::Marker mt;
                mt.header.stamp = t;
                mt.frame_locked = true;
                mt.header.frame_id = common_frame;
                mt.ns = "lnk/label/" + object_name_;
                mt.type = visualization_msgs::Marker::TEXT_VIEW_FACING;
                mt.text = link_name;
                mt.pose.position.x = p.x();
                mt.pose.position.y = p.y();
                mt.pose.position.z = p.z();
                mt.scale.x = 0.05;
                mt.scale.y = 0.05;
                mt.scale.z = 0.05;
                mt.color.r = 1;
                mt.color.g = 1;
                mt.color.b = 1;
                mt.color.a = 1;
                ma.markers.push_back(mt);
                pub_marker_label.publish(ma);
            }
        }
    }
}

void Point2Line::Instantiate(Point2LineInitializer &init)
{
    link_name = Frames[0].FrameALinkName;
    base_name = Frames[0].FrameBLinkName;

    line_start = Eigen::Map<Eigen::Vector3d>(Frames[0].FrameBOffset.p.data);
    line_end = init.EndPoint;

    line = line_end - line_start;
    infinite = init.Infinite;

    visualise = init.Visualise;

    if (visualise && Server::isRos())
    {
        pub_marker = Server::advertise<visualization_msgs::MarkerArray>("p2l", 1, true);
        pub_marker_label = Server::advertise<visualization_msgs::MarkerArray>("p2l_label", 1, true);
        // delete previous markers
        visualization_msgs::Marker md;
        md.action = 3;  // DELETEALL
        visualization_msgs::MarkerArray ma;
        ma.markers.push_back(md);
        pub_marker.publish(ma);
        pub_marker_label.publish(ma);
    }
}

int Point2Line::taskSpaceDim()
{
    return Kinematics[0].Phi.rows() * 3;
}
}  // namespace exotica
<commit_msg>correct gradient for 0-clipped distance (distance to line start)<commit_after>/*
 *      Author: Christian Rauch
 *
 * Copyright (c) 2018, University Of Edinburgh
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *  * Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of  nor the names of its contributors may be used to
 *    endorse or promote products derived from this software without specific
 *    prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 */

#include "task_map/Point2Line.hpp"

REGISTER_TASKMAP_TYPE("Point2Line", exotica::Point2Line);

namespace exotica
{
Eigen::Vector3d Point2Line::direction(const Eigen::Vector3d &point)
{
    // http://mathworld.wolfram.com/Point-LineDistance3-Dimensional.html
    // let:
    //      s: start of line
    //      e: end of line
    //      p: point
    // then the point on vector v = e-s that is closest to p is vp = s + t*(e-s) with
    // t = -((s-p)*(e-s)) / (|e-s|^2)    (* denotes the dot product)
    if (debug_) HIGHLIGHT_NAMED("P2L", "\e[4m" << link_name << "\e[0m");
    if (debug_) HIGHLIGHT_NAMED("P2L", "p " << point.transpose());
    if (debug_) HIGHLIGHT_NAMED("P2L", "ls " << line_start.transpose());
    if (debug_) HIGHLIGHT_NAMED("P2L", "le " << line_end.transpose());
    double t = -(line_start - point).dot(line) / line.squaredNorm();
    std::stringstream ss;
    ss << "t " << t;
    if (!infinite)
    {
        // clip to to range [0,1]
        t = std::min(std::max(0.0, t), 1.0);
        ss << ", clipped |t| " << t;
    }
    if (debug_) HIGHLIGHT_NAMED("P2L", ss.str());

    // vector from point 'p' to point 'vp' on line
    // vp = line_start + t * (line_end-line_start)
    const Eigen::Vector3d dv = line_start + (t * line) - point;
    if (debug_) HIGHLIGHT_NAMED("P2L", "vp " << (line_start + (t * line)).transpose());
    if (debug_) HIGHLIGHT_NAMED("P2L", "dv " << dv.transpose());
    return dv;
}

void Point2Line::update(Eigen::VectorXdRefConst x, Eigen::VectorXdRef phi)
{
    if (phi.rows() != Kinematics[0].Phi.rows() * 3) throw_named("Wrong size of phi!");

    for (int i = 0; i < Kinematics[0].Phi.rows(); i++)
    {
        const Eigen::Vector3d p = line_start + Eigen::Map<const Eigen::Vector3d>(Kinematics[0].Phi(i).p.data);
        phi.segment<3>(i * 3) = -direction(p);
    }
}

void Point2Line::update(Eigen::VectorXdRefConst x, Eigen::VectorXdRef phi, Eigen::MatrixXdRef J)
{
    if (phi.rows() != Kinematics[0].Phi.rows() * 3) throw_named("Wrong size of phi!");
    if (J.rows() != Kinematics[0].J.rows() * 3 || J.cols() != Kinematics[0].J(0).data.cols()) throw_named("Wrong size of J! " << Kinematics[0].J(0).data.cols());

    for (int i = 0; i < Kinematics[0].Phi.rows(); i++)
    {
        // point in base frame
        const Eigen::Vector3d p = line_start + Eigen::Map<const Eigen::Vector3d>(Kinematics[0].Phi(i).p.data);
        // direction from point to line
        const Eigen::Vector3d dv = direction(p);
        phi.segment<3>(i * 3) = dv;

        if ((dv + p - line_start).norm() < std::numeric_limits<double>::epsilon())
        {
            // clipped (t=0) case
            J.middleRows<3>(i * 3) = -Kinematics[0].J[i].data.topRows<3>();
        }
        else
        {
            for (int j = 0; j < J.cols(); j++)
            {
                J.middleRows<3>(i * 3).col(j) = Kinematics[0].J[i].data.topRows<3>().col(j).dot(line / line.squaredNorm()) * line - Kinematics[0].J[i].data.topRows<3>().col(j);
            }
        }

        // visualisation of point, line and their distance
        if (visualise && Server::isRos())
        {
            const ros::Time t = ros::Time::now();
            const std::string common_frame = "exotica/" + base_name;
            visualization_msgs::MarkerArray ma;
            {
                // line in base frame
                visualization_msgs::Marker mc;
                mc.header.stamp = t;
                mc.frame_locked = true;
                mc.header.frame_id = common_frame;
                mc.ns = "cam/line/" + object_name_;
                mc.type = visualization_msgs::Marker::ARROW;
                mc.scale.x = 0.01;
                mc.scale.y = 0.01;
                mc.scale.z = 0.01;
                // line start
                geometry_msgs::Point pp;
                pp.x = line_start.x();
                pp.y = line_start.y();
                pp.z = line_start.z();
                mc.points.push_back(pp);
                // line end
                const Eigen::Vector3d pe = p + dv;
                pp.x = pe.x();
                pp.y = pe.y();
                pp.z = pe.z();
                mc.points.push_back(pp);
                mc.color.r = 1;
                mc.color.g = 1;
                mc.color.b = 0;
                mc.color.a = 1;
                ma.markers.push_back(mc);
            }
            {
                // point in link frame
                visualization_msgs::Marker ml;
                ml.header.stamp = t;
                ml.frame_locked = true;
                ml.header.frame_id = common_frame;
                ml.ns = "lnk/point/" + object_name_;
                ml.type = visualization_msgs::Marker::SPHERE;
                ml.scale.x = 0.03;
                ml.scale.y = 0.03;
                ml.scale.z = 0.03;
                ml.color.r = 1;
                ml.color.g = 0;
                ml.color.b = 0;
                ml.color.a = 1;
                ml.pose.position.x = p.x();
                ml.pose.position.y = p.y();
                ml.pose.position.z = p.z();
                ma.markers.push_back(ml);
            }
            {
                // draw 'dv' starting at 'p' in base frame
                visualization_msgs::Marker mdv;
                mdv.header.stamp = t;
                mdv.frame_locked = true;
                mdv.header.frame_id = common_frame;
                mdv.ns = "dv/" + object_name_;
                mdv.type = visualization_msgs::Marker::ARROW;
                mdv.scale.x = 0.001;
                mdv.scale.y = 0.01;
                mdv.scale.z = 0.01;
                mdv.pose.position.x = p.x();
                mdv.pose.position.y = p.y();
                mdv.pose.position.z = p.z();
                mdv.points.push_back(geometry_msgs::Point());
                geometry_msgs::Point pdv;
                pdv.x = dv.x();
                pdv.y = dv.y();
                pdv.z = dv.z();
                mdv.points.push_back(pdv);
                mdv.color.r = 0;
                mdv.color.g = 1;
                mdv.color.b = 0;
                mdv.color.a = 0.5;
                ma.markers.push_back(mdv);
            }
            pub_marker.publish(ma);
            {
                visualization_msgs::MarkerArray ma;
                visualization_msgs::Marker mt;
                mt.header.stamp = t;
                mt.frame_locked = true;
                mt.header.frame_id = common_frame;
                mt.ns = "lnk/label/" + object_name_;
                mt.type = visualization_msgs::Marker::TEXT_VIEW_FACING;
                mt.text = link_name;
                mt.pose.position.x = p.x();
                mt.pose.position.y = p.y();
                mt.pose.position.z = p.z();
                mt.scale.x = 0.05;
                mt.scale.y = 0.05;
                mt.scale.z = 0.05;
                mt.color.r = 1;
                mt.color.g = 1;
                mt.color.b = 1;
                mt.color.a = 1;
                ma.markers.push_back(mt);
                pub_marker_label.publish(ma);
            }
        }
    }
}

void Point2Line::Instantiate(Point2LineInitializer &init)
{
    link_name = Frames[0].FrameALinkName;
    base_name = Frames[0].FrameBLinkName;

    line_start = Eigen::Map<Eigen::Vector3d>(Frames[0].FrameBOffset.p.data);
    line_end = init.EndPoint;

    line = line_end - line_start;
    infinite = init.Infinite;

    visualise = init.Visualise;

    if (visualise && Server::isRos())
    {
        pub_marker = Server::advertise<visualization_msgs::MarkerArray>("p2l", 1, true);
        pub_marker_label = Server::advertise<visualization_msgs::MarkerArray>("p2l_label", 1, true);
        // delete previous markers
        visualization_msgs::Marker md;
        md.action = 3;  // DELETEALL
        visualization_msgs::MarkerArray ma;
        ma.markers.push_back(md);
        pub_marker.publish(ma);
        pub_marker_label.publish(ma);
    }
}

int Point2Line::taskSpaceDim()
{
    return Kinematics[0].Phi.rows() * 3;
}
}  // namespace exotica
<|endoftext|>