<commit_before>/*
 * Jamoma Shared Library
 * Functions and resources used by Jamoma objects.
 * By Tim Place, Copyright © 2007
 *
 * License: This code is licensed under the terms of the "New BSD License"
 * http://creativecommons.org/licenses/BSD/
 */

#include "Jamoma.h"
#include "JamomaObject.h"

#include "AsyncRamp.h"
#include "NoneRamp.h"
#include "QueueRamp.h"
#include "SchedulerRamp.h"

// constants
const double k_pi = 3.1415926535897932;		// pi
const double k_twopi = 6.2831853071795864;	// 2 * pi
const double k_anti_denormal_value = 1e-18;

// statics and globals
static long			initialized = false;
static t_hashtab	*hash_modules = NULL;			// a hashtab of all modules (jcom.hubs) currently instantiated
//t_object			*obj_jamoma_clock = NULL;		// there is only one master clock for the whole system
//t_object			*obj_jamoma_scheduler = NULL;	// a shared global instance of the scheduler class (there may be others too)
bool				max5 = false;

/************************************************************************************/
// Init the framework

void jamoma_init(void)
{
	if (!initialized) {
		t_object*	max = SymbolGen("max")->s_thing;
		t_symbol*	meth = SymbolGen("objectfile");
		t_atom		a[4];

		if (maxversion() >= 0x0500)
			max5 = true;

		TTDSPInit();

		AsyncRamp::registerClass();
		NoneRamp::registerClass();
		QueueRamp::registerClass();
		SchedulerRamp::registerClass();

		common_symbols_init();
		jamomaSymbolsInit();

		receivemaster_initclass();
		receive_initclass();
		object_method(max, meth, SymbolGen("jcom.receive"), SymbolGen("jcom.loader"), SymbolGen("jcom.receive"));
		object_method_sym(max, SymbolGen("db.object_addinternal"), SymbolGen("jcom.receive"), NULL);

		send_initclass();
		object_method(max, meth, SymbolGen("jcom.send"), SymbolGen("jcom.loader"), SymbolGen("jcom.send"));
		object_method_sym(max, SymbolGen("db.object_addinternal"), SymbolGen("jcom.send"), NULL);

		receive_tilde_initclass();
		object_method(max, meth, SymbolGen("jcom.receive~"), SymbolGen("jcom.loader"), SymbolGen("jcom.receive~"));
		object_method_sym(max, SymbolGen("db.object_addinternal"), SymbolGen("jcom.receive~"), NULL);

		send_tilde_initclass();
		object_method(max, meth, SymbolGen("jcom.send~"), SymbolGen("jcom.loader"), SymbolGen("jcom.send~"));
		object_method_sym(max, SymbolGen("db.object_addinternal"), SymbolGen("jcom.send~"), NULL);

		// Setup Class Aliases for TTBlue
		// object_method(max, meth, SymbolGen("jcom.limiter~"), SymbolGen("tt.limiter~"), SymbolGen("jcom.limiter~"));
		// object_method(max, meth, SymbolGen("jcom.saturation~"), SymbolGen("tt.overdrive~"), SymbolGen("jcom.saturation~"));

		// Create Required Global Instances
		// obj_jamoma_clock = (t_object*)object_new_typed(CLASS_NOBOX, SymbolGen("jamoma.clock"), 0, NULL);
		// obj_jamoma_scheduler = (t_object*)object_new_typed(CLASS_NOBOX, SymbolGen("jamoma.scheduler"), 0, NULL);
		hash_modules = (t_hashtab*)hashtab_new(0);
		// TODO: Use quittask_install() to set up a destructor for this to free it before Max exits

		// This tells Max 5.0.6 and higher that we want the patcher files to be saved such that they are sorted.
		// Having the saved this way makes our SVN diffs much more meaningful.
		object_method_long(max, SymbolGen("sortpatcherdictonsave"), 1, NULL);

		// This tells Max 4.5.7 and higher to take any posts to the Max window and also make the
		// post to the system console, which greatly aids in debugging problems and crashes
		object_method_long(max, SymbolGen("setmirrortoconsole"), 1, NULL);

		// Add Jamoma Key Commands:

		// J -- Jamoma: a new object box with "jcom." in it
		atom_setsym(a+0, SymbolGen("J"));
		atom_setsym(a+1, SymbolGen("patcher"));
		atom_setsym(a+2, SymbolGen("inserttextobj"));
		atom_setsym(a+3, SymbolGen("jcom."));
		object_method_typed(max, SymbolGen("definecommand"), 4, a, NULL);

		// M -- Module: a new object box with "jmod." in it
		atom_setsym(a+0, SymbolGen("M"));
		atom_setsym(a+1, SymbolGen("patcher"));
		atom_setsym(a+2, SymbolGen("inserttextobj"));
		atom_setsym(a+3, SymbolGen("jmod."));
		object_method_typed(max, SymbolGen("definecommand"), 4, a, NULL);

		// I -- Input: a new audio input module
		object_method_parse(max, SymbolGen("definecommand"), (char*)"I patcher insertobj bpatcher @name jmod.input~.maxpat @args /input~", NULL);
		// O -- Output: a new audio output module
		object_method_parse(max, SymbolGen("definecommand"), (char*)"O patcher insertobj bpatcher @name jmod.output~.maxpat @args /output~", NULL);

		// B -- BPatcher: a new module in a bpatcher
		object_method_parse(max, SymbolGen("definecommand"), (char*)"B patcher inserttextobj \"bpatcher @name jmod. @args myModule\"", NULL);
		// D -- Demo: a new module in a bpatcher, but with the args reverse which is handy for super-fast demos when you don't care about the OSC name
		object_method_parse(max, SymbolGen("definecommand"), (char*)"D patcher inserttextobj \"bpatcher @name jmod.\"", NULL);

		// X -- Continuous Mapper module
		object_method_parse(max, SymbolGen("definecommand"), (char*)"X patcher insertobj bpatcher @name jmod.mapperContinuous.maxpat @args /mapper", NULL);

		// !!!! --- x is defined here to work around a 'bug' in the Max Toolbox b13 ( http://code.google.com/p/maxtoolbox/downloads/list )
		object_method_parse(max, SymbolGen("definecommand"), (char*)"x patcher nothing", NULL);

		// Here bind the TTBlue environment object to the symbol "TTBlue"
		{
			t_symbol* TTBlueMaxSymbol = SymbolGen("TTBlue");

			TTBlueMaxSymbol->s_thing = 0;
			// Before we can do this we have to have a ttblue max class to receive the messages, duh...
		}

		// now the jamoma object
		{
			t_symbol* jamomaSymbol = SymbolGen("jamoma");

			jamoma_object_initclass();
			jamomaSymbol->s_thing = jamoma_object_new();
		}

		post("Jamoma %s (build %s) - www.jamoma.org", TT_MODULAR_VERSION, TT_MODULAR_REV);
		initialized = true;
	}
}

t_max_err jamoma_hub_register(t_symbol *name, t_object *hub)
{
	t_object *test = NULL;

	hashtab_lookup(hash_modules, name, &test);
	if (test)
		return MAX_ERR_GENERIC;
	else {
		hashtab_store(hash_modules, name, hub);
		return MAX_ERR_NONE;
	}
}

void jamoma_hub_remove(t_symbol *name)
{
	hashtab_chuckkey(hash_modules, name);
}

t_object* jamoma_get_hub_for_module_named(t_symbol *name)
{
	t_object *hub = NULL;

	hashtab_lookup(hash_modules, name, &hub);
	return hub;
}

void jamoma_get_all_module_names(long *numModules, t_symbol ***moduleNames)
{
	hashtab_getkeys(hash_modules, numModules, moduleNames);
}

t_object* jamoma_object_getpatcher(t_object *obj)
{
	t_object *patcher = NULL;

	object_obex_lookup(obj, SymbolGen("#P"), &patcher);
	return patcher;
}

t_symbol *jamoma_patcher_getcontext(t_object *patcher)
{
	t_object	*box = object_attr_getobj(patcher, jps_box);
	t_symbol	*objclass = NULL;

	if (box)
		objclass = object_classname(box);

	if (objclass == SymbolGen("bpatcher"))
		return objclass;
	else if (objclass == SymbolGen("newobj"))
		return SymbolGen("subpatcher");
	else
		return SymbolGen("toplevel");
}

// Don't pass memory in for this function!  It will allocate what it needs
// -- then the caller is responsible for freeing
void jamoma_patcher_getargs(t_object *patcher, long *argc, t_atom **argv)
{
	t_symbol		*context = jamoma_patcher_getcontext(patcher);
	t_object		*box = object_attr_getobj(patcher, jps_box);
	t_object		*textfield = NULL;
	char			*text = NULL;
	unsigned long	textlen = 0;

	if (context == SymbolGen("bpatcher"))
		object_attr_getvalueof(box, SymbolGen("args"), argc, argv);
	else if (context == SymbolGen("subpatcher")) {
		textfield = object_attr_getobj(box, SymbolGen("textfield"));
		object_method(textfield, SymbolGen("gettextptr"), &text, &textlen);
		atom_setparse(argc, argv, text);
	}
	else {
		*argc = 0;
		*argv = NULL;
	}
}

t_symbol* jamoma_patcher_getvarname(t_object *patcher)
{
	if (max5) {
		return _sym_nothing;
	}
	else {
		error("This version of Jamoma requires Max 5");
		return _sym_nothing;
	}
}

// When the DSP is started, we make sure the Jamoma Clock and Scheduler
// are updated if needed
void jamoma_dsp(t_object *, t_signal **sp, short *count)
{
	//jamoma_clock_dsp(NULL, sp, count);
}

// This is the default attribute getter for Jamoma's core
// It sends the data out not only the parameter dumpout, but also
// out the feedback outlet from the module
// for example:
// 	jamoma_class_attr_new(c, "ramp/drive", _sym_symbol, (method)param_setramp,
//		calcoffset(t_param, attr_ramp));
void jamoma_class_attr_new(t_class *c, const char *attrName, t_symbol *attrType, method setter, method getter)
{
	t_object	*attr = NULL;
	char		getterName[256];

	strcpy(getterName, attrName);
	strcat(getterName, "/get");

	attr = attr_offset_new(attrName, attrType, 0, getter, setter, NULL);
	class_addattr(c, attr);

	class_addmethod(c, (method)jamoma_class_attr_get, getterName, A_GIMME, 0);
}

void jamoma_class_attr_array_new(t_class *c, const char *attrName, t_symbol *attrType, long list_size, method setter, method getter)
{
	t_object	*attr = NULL;
	char		getterName[256];

	strcpy(getterName, attrName);
	strcat(getterName, "/get");

	attr = attr_offset_array_new((char*)attrName, _sym_atom, list_size, 0, getter, setter, NULL, NULL);
	class_addattr(c, attr);

	class_addmethod(c, (method)jamoma_class_attr_get, getterName, A_GIMME, 0);
}

void jamoma_class_attr_get(t_object *o, t_symbol *attrName, long, t_atom *)
{
	char		cAttrName[256];
	t_symbol	*sAttrName;
	char		*temp;
	long		ac = 0;
	t_atom		*av = NULL;
	t_jcom_core_subscriber_common *x = (t_jcom_core_subscriber_common*)o;

	strcpy(cAttrName, attrName->s_name);
	temp = strrchr(cAttrName, '/');
	if (temp)
		*temp = 0;
	sAttrName = SymbolGen(cAttrName);

	object_attr_getvalueof(o, sAttrName, &ac, &av);
	object_obex_dumpout(o, sAttrName, ac, av);
	if (x->hub != NULL) {
		char		s[256];
		t_atom		a[4];

		snprintf(s, 256, "%s:/%s", x->attr_name->s_name, attrName->s_name);
		atom_setsym(a+0, SymbolGen(s));
		sysmem_copyptr(av, a+1, sizeof(t_atom) * ac);
		object_method_typed(x->hub, jps_feedback, ac + 1, a, NULL);
	}

	if (ac)
		sysmem_freeptr(av);
}

<commit_msg>fixed crash getting an array param with more than three atoms<commit_after>/*
 * Jamoma Shared Library
 * Functions and resources used by Jamoma objects.
 * By Tim Place, Copyright © 2007
 *
 * License: This code is licensed under the terms of the "New BSD License"
 * http://creativecommons.org/licenses/BSD/
 */

#include "Jamoma.h"
#include "JamomaObject.h"

#include "AsyncRamp.h"
#include "NoneRamp.h"
#include "QueueRamp.h"
#include "SchedulerRamp.h"

// constants
const double k_pi = 3.1415926535897932;		// pi
const double k_twopi = 6.2831853071795864;	// 2 * pi
const double k_anti_denormal_value = 1e-18;

// statics and globals
static long			initialized = false;
static t_hashtab	*hash_modules = NULL;			// a hashtab of all modules (jcom.hubs) currently instantiated
//t_object			*obj_jamoma_clock = NULL;		// there is only one master clock for the whole system
//t_object			*obj_jamoma_scheduler = NULL;	// a shared global instance of the scheduler class (there may be others too)
bool				max5 = false;

/************************************************************************************/
// Init the framework

void jamoma_init(void)
{
	if (!initialized) {
		t_object*	max = SymbolGen("max")->s_thing;
		t_symbol*	meth = SymbolGen("objectfile");
		t_atom		a[4];

		if (maxversion() >= 0x0500)
			max5 = true;

		TTDSPInit();

		AsyncRamp::registerClass();
		NoneRamp::registerClass();
		QueueRamp::registerClass();
		SchedulerRamp::registerClass();

		common_symbols_init();
		jamomaSymbolsInit();

		receivemaster_initclass();
		receive_initclass();
		object_method(max, meth, SymbolGen("jcom.receive"), SymbolGen("jcom.loader"), SymbolGen("jcom.receive"));
		object_method_sym(max, SymbolGen("db.object_addinternal"), SymbolGen("jcom.receive"), NULL);

		send_initclass();
		object_method(max, meth, SymbolGen("jcom.send"), SymbolGen("jcom.loader"), SymbolGen("jcom.send"));
		object_method_sym(max, SymbolGen("db.object_addinternal"), SymbolGen("jcom.send"), NULL);

		receive_tilde_initclass();
		object_method(max, meth, SymbolGen("jcom.receive~"), SymbolGen("jcom.loader"), SymbolGen("jcom.receive~"));
		object_method_sym(max, SymbolGen("db.object_addinternal"), SymbolGen("jcom.receive~"), NULL);

		send_tilde_initclass();
		object_method(max, meth, SymbolGen("jcom.send~"), SymbolGen("jcom.loader"), SymbolGen("jcom.send~"));
		object_method_sym(max, SymbolGen("db.object_addinternal"), SymbolGen("jcom.send~"), NULL);

		// Setup Class Aliases for TTBlue
		// object_method(max, meth, SymbolGen("jcom.limiter~"), SymbolGen("tt.limiter~"), SymbolGen("jcom.limiter~"));
		// object_method(max, meth, SymbolGen("jcom.saturation~"), SymbolGen("tt.overdrive~"), SymbolGen("jcom.saturation~"));

		// Create Required Global Instances
		// obj_jamoma_clock = (t_object*)object_new_typed(CLASS_NOBOX, SymbolGen("jamoma.clock"), 0, NULL);
		// obj_jamoma_scheduler = (t_object*)object_new_typed(CLASS_NOBOX, SymbolGen("jamoma.scheduler"), 0, NULL);
		hash_modules = (t_hashtab*)hashtab_new(0);
		// TODO: Use quittask_install() to set up a destructor for this to free it before Max exits

		// This tells Max 5.0.6 and higher that we want the patcher files to be saved such that they are sorted.
		// Having the saved this way makes our SVN diffs much more meaningful.
		object_method_long(max, SymbolGen("sortpatcherdictonsave"), 1, NULL);

		// This tells Max 4.5.7 and higher to take any posts to the Max window and also make the
		// post to the system console, which greatly aids in debugging problems and crashes
		object_method_long(max, SymbolGen("setmirrortoconsole"), 1, NULL);

		// Add Jamoma Key Commands:

		// J -- Jamoma: a new object box with "jcom." in it
		atom_setsym(a+0, SymbolGen("J"));
		atom_setsym(a+1, SymbolGen("patcher"));
		atom_setsym(a+2, SymbolGen("inserttextobj"));
		atom_setsym(a+3, SymbolGen("jcom."));
		object_method_typed(max, SymbolGen("definecommand"), 4, a, NULL);

		// M -- Module: a new object box with "jmod." in it
		atom_setsym(a+0, SymbolGen("M"));
		atom_setsym(a+1, SymbolGen("patcher"));
		atom_setsym(a+2, SymbolGen("inserttextobj"));
		atom_setsym(a+3, SymbolGen("jmod."));
		object_method_typed(max, SymbolGen("definecommand"), 4, a, NULL);

		// I -- Input: a new audio input module
		object_method_parse(max, SymbolGen("definecommand"), (char*)"I patcher insertobj bpatcher @name jmod.input~.maxpat @args /input~", NULL);
		// O -- Output: a new audio output module
		object_method_parse(max, SymbolGen("definecommand"), (char*)"O patcher insertobj bpatcher @name jmod.output~.maxpat @args /output~", NULL);

		// B -- BPatcher: a new module in a bpatcher
		object_method_parse(max, SymbolGen("definecommand"), (char*)"B patcher inserttextobj \"bpatcher @name jmod. @args myModule\"", NULL);
		// D -- Demo: a new module in a bpatcher, but with the args reverse which is handy for super-fast demos when you don't care about the OSC name
		object_method_parse(max, SymbolGen("definecommand"), (char*)"D patcher inserttextobj \"bpatcher @name jmod.\"", NULL);

		// X -- Continuous Mapper module
		object_method_parse(max, SymbolGen("definecommand"), (char*)"X patcher insertobj bpatcher @name jmod.mapperContinuous.maxpat @args /mapper", NULL);

		// !!!! --- x is defined here to work around a 'bug' in the Max Toolbox b13 ( http://code.google.com/p/maxtoolbox/downloads/list )
		object_method_parse(max, SymbolGen("definecommand"), (char*)"x patcher nothing", NULL);

		// Here bind the TTBlue environment object to the symbol "TTBlue"
		{
			t_symbol* TTBlueMaxSymbol = SymbolGen("TTBlue");

			TTBlueMaxSymbol->s_thing = 0;
			// Before we can do this we have to have a ttblue max class to receive the messages, duh...
		}

		// now the jamoma object
		{
			t_symbol* jamomaSymbol = SymbolGen("jamoma");

			jamoma_object_initclass();
			jamomaSymbol->s_thing = jamoma_object_new();
		}

		post("Jamoma %s (build %s) - www.jamoma.org", TT_MODULAR_VERSION, TT_MODULAR_REV);
		initialized = true;
	}
}

t_max_err jamoma_hub_register(t_symbol *name, t_object *hub)
{
	t_object *test = NULL;

	hashtab_lookup(hash_modules, name, &test);
	if (test)
		return MAX_ERR_GENERIC;
	else {
		hashtab_store(hash_modules, name, hub);
		return MAX_ERR_NONE;
	}
}

void jamoma_hub_remove(t_symbol *name)
{
	hashtab_chuckkey(hash_modules, name);
}

t_object* jamoma_get_hub_for_module_named(t_symbol *name)
{
	t_object *hub = NULL;

	hashtab_lookup(hash_modules, name, &hub);
	return hub;
}

void jamoma_get_all_module_names(long *numModules, t_symbol ***moduleNames)
{
	hashtab_getkeys(hash_modules, numModules, moduleNames);
}

t_object* jamoma_object_getpatcher(t_object *obj)
{
	t_object *patcher = NULL;

	object_obex_lookup(obj, SymbolGen("#P"), &patcher);
	return patcher;
}

t_symbol *jamoma_patcher_getcontext(t_object *patcher)
{
	t_object	*box = object_attr_getobj(patcher, jps_box);
	t_symbol	*objclass = NULL;

	if (box)
		objclass = object_classname(box);

	if (objclass == SymbolGen("bpatcher"))
		return objclass;
	else if (objclass == SymbolGen("newobj"))
		return SymbolGen("subpatcher");
	else
		return SymbolGen("toplevel");
}

// Don't pass memory in for this function!  It will allocate what it needs
// -- then the caller is responsible for freeing
void jamoma_patcher_getargs(t_object *patcher, long *argc, t_atom **argv)
{
	t_symbol		*context = jamoma_patcher_getcontext(patcher);
	t_object		*box = object_attr_getobj(patcher, jps_box);
	t_object		*textfield = NULL;
	char			*text = NULL;
	unsigned long	textlen = 0;

	if (context == SymbolGen("bpatcher"))
		object_attr_getvalueof(box, SymbolGen("args"), argc, argv);
	else if (context == SymbolGen("subpatcher")) {
		textfield = object_attr_getobj(box, SymbolGen("textfield"));
		object_method(textfield, SymbolGen("gettextptr"), &text, &textlen);
		atom_setparse(argc, argv, text);
	}
	else {
		*argc = 0;
		*argv = NULL;
	}
}

t_symbol* jamoma_patcher_getvarname(t_object *patcher)
{
	if (max5) {
		return _sym_nothing;
	}
	else {
		error("This version of Jamoma requires Max 5");
		return _sym_nothing;
	}
}

// When the DSP is started, we make sure the Jamoma Clock and Scheduler
// are updated if needed
void jamoma_dsp(t_object *, t_signal **sp, short *count)
{
	//jamoma_clock_dsp(NULL, sp, count);
}

// This is the default attribute getter for Jamoma's core
// It sends the data out not only the parameter dumpout, but also
// out the feedback outlet from the module
// for example:
// 	jamoma_class_attr_new(c, "ramp/drive", _sym_symbol, (method)param_setramp,
//		calcoffset(t_param, attr_ramp));
void jamoma_class_attr_new(t_class *c, const char *attrName, t_symbol *attrType, method setter, method getter)
{
	t_object	*attr = NULL;
	char		getterName[256];

	strcpy(getterName, attrName);
	strcat(getterName, "/get");

	attr = attr_offset_new(attrName, attrType, 0, getter, setter, NULL);
	class_addattr(c, attr);

	class_addmethod(c, (method)jamoma_class_attr_get, getterName, A_GIMME, 0);
}

void jamoma_class_attr_array_new(t_class *c, const char *attrName, t_symbol *attrType, long list_size, method setter, method getter)
{
	t_object	*attr = NULL;
	char		getterName[256];

	strcpy(getterName, attrName);
	strcat(getterName, "/get");

	attr = attr_offset_array_new((char*)attrName, _sym_atom, list_size, 0, getter, setter, NULL, NULL);
	class_addattr(c, attr);

	class_addmethod(c, (method)jamoma_class_attr_get, getterName, A_GIMME, 0);
}

void jamoma_class_attr_get(t_object *o, t_symbol *attrName, long, t_atom *)
{
	char		cAttrName[256];
	t_symbol	*sAttrName;
	char		*temp;
	long		ac = 0;
	t_atom		*av = NULL;
	t_jcom_core_subscriber_common *x = (t_jcom_core_subscriber_common*)o;

	strcpy(cAttrName, attrName->s_name);
	temp = strrchr(cAttrName, '/');
	if (temp)
		*temp = 0;
	sAttrName = SymbolGen(cAttrName);

	object_attr_getvalueof(o, sAttrName, &ac, &av);
	object_obex_dumpout(o, sAttrName, ac, av);
	if (x->hub != NULL) {
		char		s[256];
		t_atom		*a = (t_atom *) alloca((ac + 1) * sizeof(t_atom));

		snprintf(s, 256, "%s:/%s", x->attr_name->s_name, attrName->s_name);
		atom_setsym(a+0, SymbolGen(s));
		sysmem_copyptr(av, a+1, sizeof(t_atom) * ac);
		object_method_typed(x->hub, jps_feedback, ac + 1, a, NULL);
	}

	if (ac)
		sysmem_freeptr(av);
}

<|endoftext|>


<commit_before>/*
 * Sequencer.cpp
 *
 *  Created on: Dec 24, 2018
 *      Author: david
 */

#include "Sequencer.h"

Sequencer::Sequencer() {
	_enabled = 0;

}

void Sequencer::begin(RGBOutput* rgb) {
	_enabled = 0;
	_currentBank = 0;
	_currentStep = 0;
	_markStartTime = 0;
	_lastChangeTime = micros();
	_nextChangeTime = _lastChangeTime;
	_sequenceDelay = 500000L; // .5 seconds
	_rgb=rgb;
}

void Sequencer::setChannel(int bank, int step, Channel * channel) {
	_stepChannels[bank][step] = channel;
	Log.Info("[Seq][%d][%d] is channel %d"CR,
							bank,
							step,
							channel == NULL ? -1 : channel->getNumber());
}

void Sequencer::setBank(int bank) {
	// Reset before changing banks to turn off any channels currently on
	_currentBank = bank;
	Log.Info("[Seq] Bank %d selected"CR, bank);
	reset();
}

void Sequencer::on() {
	_enabled = 1;
	Log.Info("[Seq] on"CR);
	reset();
}

void Sequencer::off() {
	_enabled = 0;
	reset();
	Log.Info("[Seq] off"CR);
}

void Sequencer::toggle() {
	if(_enabled) {
		off();
	} else {
		on();
	}
}

void Sequencer::reset() {
	_markStartTime = 0;
	_lastChangeTime = micros();
	_nextChangeTime = _lastChangeTime + _sequenceDelay;
	_currentStep = 0;

	// Turn everything off
	for(int step = 0; step < SEQ_STEPS; step++) {
		updateChannel(step, false);
	}

	if(_rgb != NULL) {
		if(_rgb->getMode() == RGBOutput::SEQUENCE_RGB || _rgb->getMode() == RGBOutput::SEQUENCE_RANDOM) {
			_rgb->setLevels(random(PIN_MIN_VALUE), random(PIN_MIN_VALUE), random(PIN_MIN_VALUE));
		}
	}
}

void Sequencer::update() {
	int redLevel;
	int greenLevel;
	int blueLevel;

	if(_enabled) {
		// Mark the current time the first thing for greatest accuracy on the next sequence time.
		unsigned long currentTime = micros();
		if(currentTime >= _nextChangeTime) {
			// Time to change, move to the next step. Turn off the current channel then advance to next
			updateChannel(_currentStep, false);
			if(_currentStep == (SEQ_STEPS - 1)) {
				_currentStep = 0;
			} else {
				_currentStep++;
			}
			updateChannel(_currentStep, true);
			Log.Info("[Seq][%d][%d] (%d)"CR, _currentBank, _currentStep, _stepChannels[_currentBank][_currentStep]->getNumber());

			_lastChangeTime = currentTime;
			_nextChangeTime = currentTime + _sequenceDelay;

			// Is RGBOutput mode SEQUENCE?
			if(_rgb != NULL) {
				if(_rgb->getMode() == RGBOutput::SEQUENCE_RGB || _rgb->getMode() == RGBOutput::SEQUENCE_RANDOM) {
					if(_rgb->getMode() == RGBOutput::SEQUENCE_RGB) {
						// @TODO Use modulo so it can handle any number of steps
						redLevel   = _currentStep == 0 ? PIN_MAX_VALUE : PIN_MIN_VALUE;
						greenLevel = _currentStep == 1 ? PIN_MAX_VALUE : PIN_MIN_VALUE;
						blueLevel  = _currentStep == 2 ? PIN_MAX_VALUE : PIN_MIN_VALUE;
					} else if(_rgb->getMode() == RGBOutput::SEQUENCE_RANDOM) {
						redLevel   = random(PIN_MAX_VALUE);
						greenLevel = random(PIN_MAX_VALUE);
						blueLevel  = random(PIN_MAX_VALUE);
					}
					Log.Info("[Seq]RGB [%d][%d][%d]"CR, redLevel, greenLevel, blueLevel);
					_rgb->setLevels(redLevel, greenLevel, blueLevel);
				}
			}
		}
	}

}

void Sequencer::markTime() {
	if(_markStartTime == 0) {
		// Set first mark time
		_markStartTime = micros();
		Log.Info("[Seq] Mark time start"CR);
	} else {
		// Calculate the difference and set the delay
		// Set last change time so we sequence the next update
		_sequenceDelay = micros() - _markStartTime;
		_markStartTime = 0;
		Log.Info("[Seq] Delay now %dms"CR, _sequenceDelay / 1000L);
		reset();
	}
}

void Sequencer::updateChannel(int step, bool state) {
	if(_stepChannels[_currentBank][step] != NULL) {
		if (state == true) {
			_stepChannels[_currentBank][step]->on();
		} else {
			_stepChannels[_currentBank][step]->off();
		}
	}
}

Sequencer::~Sequencer() {

}
<commit_msg>Reset BEFORE changing banks<commit_after>/*
 * Sequencer.cpp
 *
 *  Created on: Dec 24, 2018
 *      Author: david
 */

#include "Sequencer.h"

Sequencer::Sequencer() {
	_enabled = 0;

}

void Sequencer::begin(RGBOutput* rgb) {
	_enabled = 0;
	_currentBank = 0;
	_currentStep = 0;
	_markStartTime = 0;
	_lastChangeTime = micros();
	_nextChangeTime = _lastChangeTime;
	_sequenceDelay = 500000L; // .5 seconds
	_rgb=rgb;
}

void Sequencer::setChannel(int bank, int step, Channel * channel) {
	_stepChannels[bank][step] = channel;
	Log.Info("[Seq][%d][%d] is channel %d"CR,
							bank,
							step,
							channel == NULL ? -1 : channel->getNumber());
}

void Sequencer::setBank(int bank) {
	// Reset before changing banks to turn off any channels currently on
	reset();
	_currentBank = bank;
	Log.Info("[Seq] Bank %d selected"CR, bank);
}

void Sequencer::on() {
	_enabled = 1;
	Log.Info("[Seq] on"CR);
	reset();
}

void Sequencer::off() {
	_enabled = 0;
	reset();
	Log.Info("[Seq] off"CR);
}

void Sequencer::toggle() {
	if(_enabled) {
		off();
	} else {
		on();
	}
}

void Sequencer::reset() {
	_markStartTime = 0;
	_lastChangeTime = micros();
	_nextChangeTime = _lastChangeTime + _sequenceDelay;
	_currentStep = 0;

	// Turn everything off
	for(int step = 0; step < SEQ_STEPS; step++) {
		updateChannel(step, false);
	}

	if(_rgb != NULL) {
		if(_rgb->getMode() == RGBOutput::SEQUENCE_RGB || _rgb->getMode() == RGBOutput::SEQUENCE_RANDOM) {
			_rgb->setLevels(random(PIN_MIN_VALUE), random(PIN_MIN_VALUE), random(PIN_MIN_VALUE));
		}
	}
}

void Sequencer::update() {
	int redLevel;
	int greenLevel;
	int blueLevel;

	if(_enabled) {
		// Mark the current time the first thing for greatest accuracy on the next sequence time.
		unsigned long currentTime = micros();
		if(currentTime >= _nextChangeTime) {
			// Time to change, move to the next step. Turn off the current channel then advance to next
			updateChannel(_currentStep, false);
			if(_currentStep == (SEQ_STEPS - 1)) {
				_currentStep = 0;
			} else {
				_currentStep++;
			}
			updateChannel(_currentStep, true);
			Log.Info("[Seq][%d][%d] (%d)"CR, _currentBank, _currentStep, _stepChannels[_currentBank][_currentStep]->getNumber());

			_lastChangeTime = currentTime;
			_nextChangeTime = currentTime + _sequenceDelay;

			// Is RGBOutput mode SEQUENCE?
			if(_rgb != NULL) {
				if(_rgb->getMode() == RGBOutput::SEQUENCE_RGB || _rgb->getMode() == RGBOutput::SEQUENCE_RANDOM) {
					if(_rgb->getMode() == RGBOutput::SEQUENCE_RGB) {
						// @TODO Use modulo so it can handle any number of steps
						redLevel   = _currentStep == 0 ? PIN_MAX_VALUE : PIN_MIN_VALUE;
						greenLevel = _currentStep == 1 ? PIN_MAX_VALUE : PIN_MIN_VALUE;
						blueLevel  = _currentStep == 2 ? PIN_MAX_VALUE : PIN_MIN_VALUE;
					} else if(_rgb->getMode() == RGBOutput::SEQUENCE_RANDOM) {
						redLevel   = random(PIN_MAX_VALUE);
						greenLevel = random(PIN_MAX_VALUE);
						blueLevel  = random(PIN_MAX_VALUE);
					}
					Log.Info("[Seq]RGB [%d][%d][%d]"CR, redLevel, greenLevel, blueLevel);
					_rgb->setLevels(redLevel, greenLevel, blueLevel);
				}
			}
		}
	}

}

void Sequencer::markTime() {
	if(_markStartTime == 0) {
		// Set first mark time
		_markStartTime = micros();
		Log.Info("[Seq] Mark time start"CR);
	} else {
		// Calculate the difference and set the delay
		// Set last change time so we sequence the next update
		_sequenceDelay = micros() - _markStartTime;
		_markStartTime = 0;
		Log.Info("[Seq] Delay now %dms"CR, _sequenceDelay / 1000L);
		reset();
	}
}

void Sequencer::updateChannel(int step, bool state) {
	if(_stepChannels[_currentBank][step] != NULL) {
		if (state == true) {
			_stepChannels[_currentBank][step]->on();
		} else {
			_stepChannels[_currentBank][step]->off();
		}
	}
}

Sequencer::~Sequencer() {

}
<|endoftext|>


<commit_before>#include "controller_node.h"
#include <std_msgs/Bool.h>
#include <std_srvs/Empty.h>
#include <iostream>

Controller::Controller() :
    speedX_(0.05f), speedY_(0.05f), speedZ_(0.05f), speedYaw_(0.01f), goal_reached_(true),
    search_for_plane_(false), stick_to_plane_(false), sticking_distance_(0.5f), correction_speed_(0.5f)
{
    std::cout << "Controller node started..." << std::endl;

    // get Ros parameters and set variables
    nh_.param("speedForward", speedX_, speedX_);
    nh_.param("speedRight", speedY_, speedY_);
    nh_.param("speedUp", speedZ_, speedZ_);
    nh_.param("speedYaw", speedYaw_, speedYaw_);

    // set subscriber and publisher
    sub_joy_ = nh_.subscribe<sensor_msgs::Joy>("joy", 10, &Controller::callback, this);
    sub_mocap_pose_ = nh_.subscribe<geometry_msgs::PoseWithCovarianceStamped>("vrpn_client/estimated_transform", 10, &Controller::setMocapPose, this);
    sub_plane_tf_ = nh_.subscribe<geometry_msgs::TransformStamped>("plane", 10, &Controller::processPlaneMsg, this);
    pub_pose_ = nh_.advertise<geometry_msgs::PoseStamped>("command/pose", 1);
    pub_stickToSurface_ = nh_.advertise<std_msgs::Bool>("stickToSurface", 10);

    std::cout << "speeds: " << speedX_ << " " << speedY_ << " " << speedZ_ << " " << speedYaw_ << std::endl;

    // identity rotation matrix as quaternion
    tf::Quaternion q;
    q.setRPY(0,0,0);

    // init mocap_tf
    mavToWorld_.setOrigin( tf::Vector3(0,0,0) );
    mavToWorld_.setRotation(q);

    // init transform
    transform_.setOrigin( tf::Vector3(0,0,0) );
    transform_.setRotation(q);

    // TBD if we should be able to change this via launch file
    // init hover transform
    hover_goal_tf_.setOrigin( tf::Vector3(0,0,1) );
    hover_goal_tf_.setRotation(q);

    // init plane_tf
    plane_tf_.setOrigin(tf::Vector3(2,0,0));
    tf::Quaternion plane_q;
    plane_q.setRPY(0,0,-M_PI/4.f);
    plane_tf_.setRotation(plane_q);

    // init snap_goal_tf
    snap_goal_tf_.setOrigin(tf::Vector3(0,0,0));
    snap_goal_tf_.setRotation(q);

    // hack
    tf::Transform mavToWorld;
    mavToWorld.setOrigin(tf::Vector3(0,0,0.117));
    mavToWorld.setRotation(tf::Quaternion(0,0,0,1));
    tf::Transform sensorToWorld;
    sensorToWorld.setOrigin(tf::Vector3(0.133,0,0.0605));
    sensorToWorld.setRotation(tf::Quaternion(0,0.17365,0,0.98481));
    sensorToMav_ = mavToWorld.inverse() * sensorToWorld;
}

void Controller::setMocapPose(const geometry_msgs::PoseWithCovarianceStamped::ConstPtr& msg)
{
    mavToWorld_.setOrigin( tf::Vector3(msg->pose.pose.position.x, msg->pose.pose.position.y, msg->pose.pose.position.z) );
    mavToWorld_.setRotation( tf::Quaternion(msg->pose.pose.orientation.x, msg->pose.pose.orientation.y, msg->pose.pose.orientation.z, msg->pose.pose.orientation.w) );

    if(!goal_reached_) {
        tf::Transform diff = mavToWorld_.inverse() * hover_goal_tf_;
        if(diff.getOrigin().length() < 0.05f) {
            goal_reached_ = true;
        }
        else {
            transform_ = diff;
        }
    }

    // world transform
    tf::Transform curr_transform;
    /// stick to plane
    if(stick_to_plane_) {
        testPlanes();
        tf::Vector3 diff = snap_goal_tf_.getOrigin() - mavToWorld_.getOrigin();
        curr_transform.setOrigin(mavToWorld_.getOrigin() + correction_speed_*diff);
        curr_transform.setRotation(snap_goal_tf_.getRotation());
    } else {
        curr_transform = mavToWorld_ * transform_;
    }

    // convert tf into pose and publish the pose
    tf::Vector3 desired_pos = curr_transform.getOrigin();
    tf::Quaternion desired_rot = curr_transform.getRotation();

    geometry_msgs::PoseStamped pose;
    // set header
    pose.header.stamp = ros::Time::now();
    pose.header.frame_id = "world";
    // set pose
    pose.pose.position.x = desired_pos.x();
    pose.pose.position.y = desired_pos.y();
    pose.pose.position.z = desired_pos.z();
    pose.pose.orientation.x = desired_rot.x();
    pose.pose.orientation.y = desired_rot.y();
    pose.pose.orientation.z = desired_rot.z();
    pose.pose.orientation.w = desired_rot.w();
    pub_pose_.publish(pose);

    // rviz debug
    br_tf_.sendTransform( tf::StampedTransform(curr_transform, ros::Time::now(), "world", "controller") );
}

void Controller::callback(const sensor_msgs::Joy::ConstPtr& joy)
{
    /// translation from controller axis
    float jx = speedX_ * joy->axes[PS3_AXIS_STICK_RIGHT_UPWARDS];
    float jy = speedY_ * joy->axes[PS3_AXIS_STICK_RIGHT_LEFTWARDS];
    float jz = speedZ_ * joy->axes[PS3_AXIS_STICK_LEFT_UPWARDS];
    // yaw from axis
    float jr = speedYaw_ * joy->axes[PS3_AXIS_STICK_LEFT_LEFTWARDS];

    // save only the latest relative transform in global variable transform_
    transform_.setOrigin( tf::Vector3(jx,jy,jz) );
    tf::Quaternion q;
    q.setRPY(0,0,jr);
    transform_.setRotation(q);

    /// buttons
    // tell PcMeshBuilder to search for a plane
    if(joy->buttons[PS3_BUTTON_REAR_RIGHT_1]) {
        search_for_plane_ = true;
        std_msgs::Bool sticking;
        sticking.data = search_for_plane_;
        pub_stickToSurface_.publish(sticking);
    }
    if(joy->buttons[PS3_BUTTON_REAR_LEFT_1]) {
        search_for_plane_ = false;
        std_msgs::Bool sticking;
        sticking.data = search_for_plane_;
        pub_stickToSurface_.publish(sticking);
    }
    // enable or disable sticking to the plane
    if(joy->buttons[PS3_BUTTON_REAR_RIGHT_2]) {
        stick_to_plane_ = true;

        // set current distance to surface as sticking distance
        Eigen::Vector3f plane_pos = Eigen::Vector3f( plane_tf_.getOrigin().x(), plane_tf_.getOrigin().y(), plane_tf_.getOrigin().z() );
        tf::Quaternion plane_q = plane_tf_.getRotation();
        Eigen::Quaternionf plane_rot = Eigen::Quaternionf(plane_q.w(),plane_q.x(),plane_q.y(),plane_q.z());
        Eigen::Vector3f normal = plane_rot*forward;
        normal.normalize();
        Eigen::Vector3f curr_pos = Eigen::Vector3f(mavToWorld_.getOrigin().x(), mavToWorld_.getOrigin().y(), mavToWorld_.getOrigin().z());
        sticking_distance_ = normal.dot(curr_pos-plane_pos);
    }
    if(joy->buttons[PS3_BUTTON_REAR_LEFT_2]) {
        stick_to_plane_ = false;
    }
    // sticking distance
    if(joy->buttons[PS3_BUTTON_CROSS_UP] && sticking_distance_ > 0.5f) {
        sticking_distance_ -= 0.005f;
    }
    if(joy->buttons[PS3_BUTTON_CROSS_DOWN]) {
        sticking_distance_ += 0.005f;
    }
}

void Controller::takeoffAndHover()
{
    std_srvs::Empty::Request request;
    std_srvs::Empty::Response response;
    ros::service::call("euroc2/takeoff", request, response);

    tf::Vector3 desired_pos = tf::Vector3( mavToWorld_.getOrigin().x(), mavToWorld_.getOrigin().y(), 1.0f);
    tf::Quaternion desired_rot = mavToWorld_.getRotation();

    hover_goal_tf_.setOrigin(desired_pos);
    hover_goal_tf_.setRotation(desired_rot);
}

void Controller::processPlaneMsg(const geometry_msgs::TransformStamped::ConstPtr& msg)
{
    // realtive plane transform
    // planeToCam = plane_tf
    plane_tf_.setOrigin( tf::Vector3(msg->transform.translation.x, msg->transform.translation.y, msg->transform.translation.z) );
    plane_tf_.setRotation( tf::Quaternion(msg->transform.rotation.x, msg->transform.rotation.y, msg->transform.rotation.z, msg->transform.rotation.w) );

    // plane forward is z should be x, conversion of systems
    // camToSensor
    tf::Quaternion q1; q1.setRPY(0,M_PI/2.0,0);
    tf::Quaternion q2; q2.setRPY(-M_PI/2.0,0,0);
    tf::Quaternion opticalToSensor = q2*q1;
    plane_tf_.setOrigin( tf::Matrix3x3(opticalToSensor)*plane_tf_.getOrigin() );
    plane_tf_.setRotation( opticalToSensor*plane_tf_.getRotation() );
    br_tf_.sendTransform( tf::StampedTransform(plane_tf_, ros::Time::now(), "euroc_hex/vi_sensor/ground_truth", "plane_Sensor") );

    // correct the relative camera offset
    // planeToMav = plane_tf
    plane_tf_ = sensorToMav_*plane_tf_;

    // transform into global coordinates
    // planeToWorld = plane_tf_
    plane_tf_ = mavToWorld_*plane_tf_;

    // rviz debug
    br_tf_.sendTransform( tf::StampedTransform(plane_tf_, ros::Time::now(), "world", "plane_Global") );
}

void Controller::testPlanes()
{
    //// mav
    // predicted mav tf
    tf::Transform mav_tf = mavToWorld_ * transform_;

    // Eigen conversions
    Eigen::Vector3f mav_pos = Eigen::Vector3f( mav_tf.getOrigin().x(),
                                               mav_tf.getOrigin().y(),
                                               mav_tf.getOrigin().z());

    //// plane
    Eigen::Vector3f plane_pos = Eigen::Vector3f( plane_tf_.getOrigin().x(), plane_tf_.getOrigin().y(), plane_tf_.getOrigin().z() );
    tf::Quaternion plane_q = plane_tf_.getRotation();
    Eigen::Quaternionf plane_rot = Eigen::Quaternionf(plane_q.w(),plane_q.x(),plane_q.y(),plane_q.z());

    //// calculations
    Eigen::Vector3f normal = plane_rot*forward;
    normal.normalize();
    // determine if mav is in front or behind plane normal, take current pos to avoid switching of sides by wrong predictions in mav_tf
    Eigen::Vector3f curr_pos = Eigen::Vector3f(mavToWorld_.getOrigin().x(), mavToWorld_.getOrigin().y(), mavToWorld_.getOrigin().z());

    int facing = normal.dot(curr_pos-plane_pos) >= 0 ? 1 : -1;
    // calculate projected position of the mav onto the plane
    Eigen::Vector3f proj_pos = mav_pos + (facing*sticking_distance_ - normal.dot(mav_pos-plane_pos))*normal;
}

int main(int argc, char** argv)
{
    ros::init(argc, argv, "te_joy_controller");
    Controller rc;
    ros::spin();
}
<commit_msg>Fixed a bug that snapping goal was never set.<commit_after>#include "controller_node.h"
#include <std_msgs/Bool.h>
#include <std_srvs/Empty.h>
#include <iostream>

Controller::Controller() :
    speedX_(1.05f), speedY_(1.05f), speedZ_(1.05f), speedYaw_(1.01f), goal_reached_(true),
    search_for_plane_(false), stick_to_plane_(false), sticking_distance_(0.5f), correction_speed_(0.5f)
{
    std::cout << "Controller node started..." << std::endl;

    // get Ros parameters and set variables
    nh_.param("speedForward", speedX_, speedX_);
    nh_.param("speedRight", speedY_, speedY_);
    nh_.param("speedUp", speedZ_, speedZ_);
    nh_.param("speedYaw", speedYaw_, speedYaw_);

    // set subscriber and publisher
    sub_joy_ = nh_.subscribe<sensor_msgs::Joy>("joy", 10, &Controller::callback, this);
    sub_mocap_pose_ = nh_.subscribe<geometry_msgs::PoseWithCovarianceStamped>("vrpn_client/estimated_transform", 10, &Controller::setMocapPose, this);
    sub_plane_tf_ = nh_.subscribe<geometry_msgs::TransformStamped>("plane", 10, &Controller::processPlaneMsg, this);
    pub_pose_ = nh_.advertise<geometry_msgs::PoseStamped>("command/pose", 1);
    pub_stickToSurface_ = nh_.advertise<std_msgs::Bool>("stickToSurface", 10);

    std::cout << "speeds: " << speedX_ << " " << speedY_ << " " << speedZ_ << " " << speedYaw_ << std::endl;

    // identity rotation matrix as quaternion
    tf::Quaternion q;
    q.setRPY(0,0,0);

    // init mocap_tf
    mavToWorld_.setOrigin( tf::Vector3(0,0,0) );
    mavToWorld_.setRotation(q);

    // init transform
    transform_.setOrigin( tf::Vector3(0,0,0) );
    transform_.setRotation(q);

    // TBD if we should be able to change this via launch file
    // init hover transform
    hover_goal_tf_.setOrigin( tf::Vector3(0,0,1) );
    hover_goal_tf_.setRotation(q);

    // init plane_tf
    plane_tf_.setOrigin(tf::Vector3(2,0,0));
    tf::Quaternion plane_q;
    plane_q.setRPY(0,0,-M_PI/4.f);
    plane_tf_.setRotation(plane_q);

    // init snap_goal_tf
    snap_goal_tf_.setOrigin(tf::Vector3(0,0,0));
    snap_goal_tf_.setRotation(q);

    // hack
    tf::Transform mavToWorld;
    mavToWorld.setOrigin(tf::Vector3(0,0,0.117));
    mavToWorld.setRotation(tf::Quaternion(0,0,0,1));
    tf::Transform sensorToWorld;
    sensorToWorld.setOrigin(tf::Vector3(0.133,0,0.0605));
    sensorToWorld.setRotation(tf::Quaternion(0,0.17365,0,0.98481));
    sensorToMav_ = mavToWorld.inverse() * sensorToWorld;
}

void Controller::setMocapPose(const geometry_msgs::PoseWithCovarianceStamped::ConstPtr& msg)
{
    mavToWorld_.setOrigin( tf::Vector3(msg->pose.pose.position.x, msg->pose.pose.position.y, msg->pose.pose.position.z) );
    mavToWorld_.setRotation( tf::Quaternion(msg->pose.pose.orientation.x, msg->pose.pose.orientation.y, msg->pose.pose.orientation.z, msg->pose.pose.orientation.w) );

    if(!goal_reached_) {
        tf::Transform diff = mavToWorld_.inverse() * hover_goal_tf_;
        if(diff.getOrigin().length() < 0.05f) {
            goal_reached_ = true;
        }
        else {
            transform_ = diff;
        }
    }

    // world transform
    tf::Transform curr_transform;
    /// stick to plane
    if(stick_to_plane_) {
        testPlanes();
        tf::Vector3 diff = snap_goal_tf_.getOrigin() - mavToWorld_.getOrigin();
        curr_transform.setOrigin(mavToWorld_.getOrigin() + correction_speed_*diff);
        curr_transform.setRotation(snap_goal_tf_.getRotation());
    } else {
        curr_transform = mavToWorld_ * transform_;
    }

    // convert tf into pose and publish the pose
    tf::Vector3 desired_pos = curr_transform.getOrigin();
    tf::Quaternion desired_rot = curr_transform.getRotation();

    geometry_msgs::PoseStamped pose;
    // set header
    pose.header.stamp = ros::Time::now();
    pose.header.frame_id = "world";
    // set pose
    pose.pose.position.x = desired_pos.x();
    pose.pose.position.y = desired_pos.y();
    pose.pose.position.z = desired_pos.z();
    pose.pose.orientation.x = desired_rot.x();
    pose.pose.orientation.y = desired_rot.y();
    pose.pose.orientation.z = desired_rot.z();
    pose.pose.orientation.w = desired_rot.w();
    pub_pose_.publish(pose);

    // rviz debug
    br_tf_.sendTransform( tf::StampedTransform(curr_transform, ros::Time::now(), "world", "controller") );
}

void Controller::callback(const sensor_msgs::Joy::ConstPtr& joy)
{
    /// translation from controller axis
    float jx = speedX_ * joy->axes[PS3_AXIS_STICK_RIGHT_UPWARDS];
    float jy = speedY_ * joy->axes[PS3_AXIS_STICK_RIGHT_LEFTWARDS];
    float jz = speedZ_ * joy->axes[PS3_AXIS_STICK_LEFT_UPWARDS];
    // yaw from axis
    float jr = speedYaw_ * joy->axes[PS3_AXIS_STICK_LEFT_LEFTWARDS];

    // save only the latest relative transform in global variable transform_
    transform_.setOrigin( tf::Vector3(jx,jy,jz) );
    tf::Quaternion q;
    q.setRPY(0,0,jr);
    transform_.setRotation(q);

    /// buttons
    // tell PcMeshBuilder to search for a plane
    if(joy->buttons[PS3_BUTTON_REAR_RIGHT_1]) {
        search_for_plane_ = true;
        std_msgs::Bool sticking;
        sticking.data = search_for_plane_;
        pub_stickToSurface_.publish(sticking);
    }
    if(joy->buttons[PS3_BUTTON_REAR_LEFT_1]) {
        search_for_plane_ = false;
        std_msgs::Bool sticking;
        sticking.data = search_for_plane_;
        pub_stickToSurface_.publish(sticking);
    }
    // enable or disable sticking to the plane
    if(joy->buttons[PS3_BUTTON_REAR_RIGHT_2]) {
        stick_to_plane_ = true;

        // set current distance to surface as sticking distance
        Eigen::Vector3f plane_pos = Eigen::Vector3f( plane_tf_.getOrigin().x(), plane_tf_.getOrigin().y(), plane_tf_.getOrigin().z() );
        tf::Quaternion plane_q = plane_tf_.getRotation();
        Eigen::Quaternionf plane_rot = Eigen::Quaternionf(plane_q.w(),plane_q.x(),plane_q.y(),plane_q.z());
        Eigen::Vector3f normal = plane_rot*forward;
        normal.normalize();
        Eigen::Vector3f curr_pos = Eigen::Vector3f(mavToWorld_.getOrigin().x(), mavToWorld_.getOrigin().y(), mavToWorld_.getOrigin().z());
        sticking_distance_ = normal.dot(curr_pos-plane_pos);
    }
    if(joy->buttons[PS3_BUTTON_REAR_LEFT_2]) {
        stick_to_plane_ = false;
    }
    // sticking distance
    if(joy->buttons[PS3_BUTTON_CROSS_UP] && sticking_distance_ > 0.5f) {
        sticking_distance_ -= 0.005f;
    }
    if(joy->buttons[PS3_BUTTON_CROSS_DOWN]) {
        sticking_distance_ += 0.005f;
    }
}

void Controller::takeoffAndHover()
{
    std_srvs::Empty::Request request;
    std_srvs::Empty::Response response;
    ros::service::call("euroc2/takeoff", request, response);

    tf::Vector3 desired_pos = tf::Vector3( mavToWorld_.getOrigin().x(), mavToWorld_.getOrigin().y(), 1.0f);
    tf::Quaternion desired_rot = mavToWorld_.getRotation();

    hover_goal_tf_.setOrigin(desired_pos);
    hover_goal_tf_.setRotation(desired_rot);
}

void Controller::processPlaneMsg(const geometry_msgs::TransformStamped::ConstPtr& msg)
{
    // realtive plane transform
    // planeToCam = plane_tf
    plane_tf_.setOrigin( tf::Vector3(msg->transform.translation.x, msg->transform.translation.y, msg->transform.translation.z) );
    plane_tf_.setRotation( tf::Quaternion(msg->transform.rotation.x, msg->transform.rotation.y, msg->transform.rotation.z, msg->transform.rotation.w) );

    // plane forward is z should be x, conversion of systems
    // camToSensor
    tf::Quaternion q1; q1.setRPY(0,M_PI/2.0,0);
    tf::Quaternion q2; q2.setRPY(-M_PI/2.0,0,0);
    tf::Quaternion opticalToSensor = q2*q1;
    plane_tf_.setOrigin( tf::Matrix3x3(opticalToSensor)*plane_tf_.getOrigin() );
    plane_tf_.setRotation( opticalToSensor*plane_tf_.getRotation() );
    br_tf_.sendTransform( tf::StampedTransform(plane_tf_, ros::Time::now(), "euroc_hex/vi_sensor/ground_truth", "plane_Sensor") );

    // correct the relative camera offset
    // planeToMav = plane_tf
    plane_tf_ = sensorToMav_*plane_tf_;

    // transform into global coordinates
    // planeToWorld = plane_tf_
    plane_tf_ = mavToWorld_*plane_tf_;

    // rviz debug
    br_tf_.sendTransform( tf::StampedTransform(plane_tf_, ros::Time::now(), "world", "plane_Global") );
}

void Controller::testPlanes()
{
    //// mav
    // predicted mav tf
    tf::Transform mav_tf = mavToWorld_ * transform_;

    // Eigen conversions
    Eigen::Vector3f mav_pos = Eigen::Vector3f( mav_tf.getOrigin().x(),
                                               mav_tf.getOrigin().y(),
                                               mav_tf.getOrigin().z());

    //// plane
    Eigen::Vector3f plane_pos = Eigen::Vector3f( plane_tf_.getOrigin().x(), plane_tf_.getOrigin().y(), plane_tf_.getOrigin().z() );
    tf::Quaternion plane_q = plane_tf_.getRotation();
    Eigen::Quaternionf plane_rot = Eigen::Quaternionf(plane_q.w(),plane_q.x(),plane_q.y(),plane_q.z());

    //// calculations
    Eigen::Vector3f normal = plane_rot*forward;
    normal.normalize();
    // determine if mav is in front or behind plane normal, take current pos to avoid switching of sides by wrong predictions in mav_tf
    Eigen::Vector3f curr_pos = Eigen::Vector3f(mavToWorld_.getOrigin().x(), mavToWorld_.getOrigin().y(), mavToWorld_.getOrigin().z());

    int facing = normal.dot(curr_pos-plane_pos) >= 0 ? 1 : -1;
    // calculate projected position of the mav onto the plane
    Eigen::Vector3f proj_pos = mav_pos + (facing*sticking_distance_ - normal.dot(mav_pos-plane_pos))*normal;

    // set snapping goal tf
    snap_goal_tf_.setOrigin( tf::Vector3(proj_pos.x(), proj_pos.y(), proj_pos.z()) );
    snap_goal_tf_.setRotation(plane_tf_.getRotation());
}

int main(int argc, char** argv)
{
    ros::init(argc, argv, "te_joy_controller");
    Controller rc;
    ros::spin();
}
<|endoftext|>


<commit_before>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#define EIGEN_USE_THREADS

#include <algorithm>
#include <numeric>
#include <unordered_map>
#include <utility>
#include <vector>

#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_util.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
#include "tensorflow/core/util/sparse/sparse_tensor.h"

namespace tensorflow {

using CPUDevice = Eigen::ThreadPoolDevice;

template <typename T>
class SparseFillEmptyRowsOp : public OpKernel {
 public:
  explicit SparseFillEmptyRowsOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const int kIndicesInput = 0;
    const int kValuesInput = 1;
    const int kDenseShapeInput = 2;
    const int kDefaultValueInput = 3;

    const int kOutputIndicesOutput = 0;
    const int kOutputValuesOutput = 1;
    const int kEmptyRowIndicatorOutput = 2;
    const int kReverseIndexMapOutput = 3;

    const Tensor& indices_t = context->input(kIndicesInput);
    const Tensor& values_t = context->input(kValuesInput);
    const Tensor& dense_shape_t = context->input(kDenseShapeInput);
    const Tensor& default_value_t = context->input(kDefaultValueInput);

    OP_REQUIRES(context, TensorShapeUtils::IsVector(dense_shape_t.shape()),
                errors::InvalidArgument("dense_shape must be a vector, saw: ",
                                        dense_shape_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsMatrix(indices_t.shape()),
                errors::InvalidArgument("indices must be a matrix, saw: ",
                                        indices_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsVector(values_t.shape()),
                errors::InvalidArgument("values must be a vector, saw: ",
                                        values_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(default_value_t.shape()),
                errors::InvalidArgument("default_value must be a scalar, saw: ",
                                        default_value_t.shape().DebugString()));
    // TODO(ebrevdo): add shape checks between values, indices,
    // dense_shape.  Also add check that dense rank > 0.

    const T& default_value = default_value_t.scalar<T>()();
    const auto indices = indices_t.matrix<int64>();
    const auto values = values_t.vec<T>();
    const auto dense_shape = dense_shape_t.vec<int64>();

    const int64 N = indices_t.shape().dim_size(0);
    const int64 dense_rows = dense_shape(0);

    bool* empty_row_indicator = nullptr;
    if (context->output_required(kEmptyRowIndicatorOutput)) {
      Tensor* empty_row_indicator_t = nullptr;
      OP_REQUIRES_OK(context,
                     context->allocate_output(kEmptyRowIndicatorOutput,
                                              TensorShape({dense_rows}),
                                              &empty_row_indicator_t));
      empty_row_indicator = empty_row_indicator_t->vec<bool>().data();
    }
    int64* reverse_index_map = nullptr;
    if (context->output_required(kReverseIndexMapOutput)) {
      Tensor* reverse_index_map_t = nullptr;
      OP_REQUIRES_OK(context, context->allocate_output(kReverseIndexMapOutput,
                                                       TensorShape({N}),
                                                       &reverse_index_map_t));
      reverse_index_map = reverse_index_map_t->vec<int64>().data();
    }

    int rank = indices_t.shape().dim_size(1);

    if (dense_rows == 0) {
      OP_REQUIRES(
          context, N == 0,
          errors::InvalidArgument("Received SparseTensor with dense_shape[0] = "
                                  "0 but indices.shape[0] = ",
                                  N));
      Tensor* output_indices_t;
      TensorShape output_indices_shape({0, rank});
      OP_REQUIRES_OK(context, context->allocate_output(kOutputIndicesOutput,
                                                       output_indices_shape,
                                                       &output_indices_t));
      Tensor* output_values_t;
      OP_REQUIRES_OK(context, context->allocate_output(kOutputValuesOutput,
                                                       TensorShape({0}),
                                                       &output_values_t));

      // Exit early, nothing more to do.
      return;
    }

    bool indices_is_order = true;
    int64 last_indices_row = 0;
    std::vector<int64> csr_offset(dense_rows, 0);
    for (int i = 0; i < N; ++i) {
      const int64 row = indices(i, 0);
      OP_REQUIRES(context, row >= 0 && row < dense_rows,
                  errors::InvalidArgument("indices(", i, ", 0) is invalid: ",
                                          row, " >= ", dense_rows));
      ++csr_offset[row];
      indices_is_order = indices_is_order & (row >= last_indices_row);
      last_indices_row = row;
    }
    bool all_rows_full = true;
    for (int row = 0; row < dense_rows; ++row) {
      // csr_offset here describes the number of elements in this dense row
      bool row_empty = (csr_offset[row] == 0);
      if (empty_row_indicator) {
        empty_row_indicator[row] = row_empty;
      }
      all_rows_full = all_rows_full & !row_empty;
      // In filled version, each row has at least one element.
      csr_offset[row] = std::max(csr_offset[row], int64{1});
      // Update csr_offset to represent the number of elements up to and
      // including dense_row + 1:
      //  csr_offset(0) == #{elements of row 0}
      //  csr_offset(1) == #{elements of row 1} + #{elements of row 0}
      //  ..
      //  csr_offset(i) == starting index for elements in row i + 1.
      if (row > 0) {
        csr_offset[row] += csr_offset[row - 1];
      }
    }

    if (all_rows_full && indices_is_order) {
      context->set_output(kOutputIndicesOutput, indices_t);
      context->set_output(kOutputValuesOutput, values_t);
      if (reverse_index_map) {
        for (int64 i = 0; i < N; ++i) {
          reverse_index_map[i] = i;
        }
      }
    } else {
      Tensor* output_indices_t;
      const int64 N_full = csr_offset[dense_rows - 1];
      TensorShape output_indices_shape({N_full, rank});
      OP_REQUIRES_OK(context, context->allocate_output(kOutputIndicesOutput,
                                                       output_indices_shape,
                                                       &output_indices_t));
      auto output_indices = output_indices_t->matrix<int64>();

      Tensor* output_values_t;
      OP_REQUIRES_OK(context, context->allocate_output(kOutputValuesOutput,
                                                       TensorShape({N_full}),
                                                       &output_values_t));
      auto output_values = output_values_t->vec<T>();

      std::vector<int64> filled_count(dense_rows, 0);

      // Fill in values for rows that are not missing
      for (int64 i = 0; i < N; ++i) {
        const int64 row = indices(i, 0);
        int64& offset = filled_count[row];
        const int64 output_i = ((row == 0) ? 0 : csr_offset[row - 1]) + offset;
        offset++;  // Increment the filled count for this row.
        std::copy_n(&indices(i, 0), rank, &output_indices(output_i, 0));
        output_values(output_i) = values(i);
        // We'll need this reverse index map to backprop correctly.
        if (reverse_index_map) {
          reverse_index_map[i] = output_i;
        }
      }

      // Fill in values for rows that are missing
      for (int64 row = 0; row < dense_rows; ++row) {
        const int64 row_count = filled_count[row];
        if (row_count == 0) {  // We haven't filled this row
          const int64 starting_index = (row == 0) ? 0 : csr_offset[row - 1];
          // Remaining index values were set to zero already.
          // Just need to set the row index in the right location.
          output_indices(starting_index, 0) = row;
          for (int64 col = 1; col < rank; ++col) {
            output_indices(starting_index, col) = 0;
          }
          output_values(starting_index) = default_value;
        }
      }
    }
  }
};

#define REGISTER_KERNELS(type)                            \
  REGISTER_KERNEL_BUILDER(Name("SparseFillEmptyRows")     \
                              .Device(DEVICE_CPU)         \
                              .TypeConstraint<type>("T"), \
                          SparseFillEmptyRowsOp<type>)

TF_CALL_ALL_TYPES(REGISTER_KERNELS);
#undef REGISTER_KERNELS

template <typename T>
class SparseFillEmptyRowsGradOp : public OpKernel {
 public:
  explicit SparseFillEmptyRowsGradOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const Tensor* reverse_index_map_t;
    const Tensor* grad_values_t;
    OP_REQUIRES_OK(context,
                   context->input("reverse_index_map", &reverse_index_map_t));
    OP_REQUIRES_OK(context, context->input("grad_values", &grad_values_t));

    const CPUDevice& d = context->eigen_device<CPUDevice>();

    OP_REQUIRES(
        context, TensorShapeUtils::IsVector(reverse_index_map_t->shape()),
        errors::InvalidArgument("reverse_index_map must be a vector, saw: ",
                                reverse_index_map_t->shape().DebugString()));

    const auto reverse_index_map = reverse_index_map_t->vec<int64>();
    const auto grad_values = grad_values_t->vec<T>();

    const int64 N = reverse_index_map_t->shape().dim_size(0);
    const int64 N_full = grad_values_t->shape().dim_size(0);

    Tensor* d_values_t;
    OP_REQUIRES_OK(context, context->allocate_output(
                                "d_values", TensorShape({N}), &d_values_t));
    auto d_values = d_values_t->vec<T>();
    Tensor* d_default_value_t;
    OP_REQUIRES_OK(context,
                   context->allocate_output("d_default_value", TensorShape({}),
                                            &d_default_value_t));
    T& d_default_value = d_default_value_t->scalar<T>()();
    d_default_value = T();

    Tensor visited_t;
    OP_REQUIRES_OK(context, context->allocate_temp(
                                DT_BOOL, TensorShape({N_full}), &visited_t));
    auto visited = visited_t.vec<bool>();
    visited.device(d) = visited.constant(false);

    for (int i = 0; i < N; ++i) {
      // Locate the index of the output of the forward prop associated
      // with this location in the input of the forward prop.  Copy
      // the gradient into it.  Mark it as visited.
      d_values(i) = grad_values(reverse_index_map(i));
      visited(reverse_index_map(i)) = true;
    }
    for (int j = 0; j < N_full; ++j) {
      // The default value gradient gets the accumulated remainder of
      // the backprop values (since the default value was used to fill
      // in these slots in the forward calculation).
      if (!visited(j)) {
        d_default_value += grad_values(j);
      }
    }
  }
};

#define REGISTER_KERNELS(type)                            \
  REGISTER_KERNEL_BUILDER(Name("SparseFillEmptyRowsGrad") \
                              .Device(DEVICE_CPU)         \
                              .TypeConstraint<type>("T"), \
                          SparseFillEmptyRowsGradOp<type>)

TF_CALL_NUMBER_TYPES(REGISTER_KERNELS);
#undef REGISTER_KERNELS
}  // namespace tensorflow
<commit_msg>Rename `indices_is_order` to `rows_are_ordered `<commit_after>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#define EIGEN_USE_THREADS

#include <algorithm>
#include <numeric>
#include <unordered_map>
#include <utility>
#include <vector>

#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_util.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
#include "tensorflow/core/util/sparse/sparse_tensor.h"

namespace tensorflow {

using CPUDevice = Eigen::ThreadPoolDevice;

template <typename T>
class SparseFillEmptyRowsOp : public OpKernel {
 public:
  explicit SparseFillEmptyRowsOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const int kIndicesInput = 0;
    const int kValuesInput = 1;
    const int kDenseShapeInput = 2;
    const int kDefaultValueInput = 3;

    const int kOutputIndicesOutput = 0;
    const int kOutputValuesOutput = 1;
    const int kEmptyRowIndicatorOutput = 2;
    const int kReverseIndexMapOutput = 3;

    const Tensor& indices_t = context->input(kIndicesInput);
    const Tensor& values_t = context->input(kValuesInput);
    const Tensor& dense_shape_t = context->input(kDenseShapeInput);
    const Tensor& default_value_t = context->input(kDefaultValueInput);

    OP_REQUIRES(context, TensorShapeUtils::IsVector(dense_shape_t.shape()),
                errors::InvalidArgument("dense_shape must be a vector, saw: ",
                                        dense_shape_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsMatrix(indices_t.shape()),
                errors::InvalidArgument("indices must be a matrix, saw: ",
                                        indices_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsVector(values_t.shape()),
                errors::InvalidArgument("values must be a vector, saw: ",
                                        values_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(default_value_t.shape()),
                errors::InvalidArgument("default_value must be a scalar, saw: ",
                                        default_value_t.shape().DebugString()));
    // TODO(ebrevdo): add shape checks between values, indices,
    // dense_shape.  Also add check that dense rank > 0.

    const T& default_value = default_value_t.scalar<T>()();
    const auto indices = indices_t.matrix<int64>();
    const auto values = values_t.vec<T>();
    const auto dense_shape = dense_shape_t.vec<int64>();

    const int64 N = indices_t.shape().dim_size(0);
    const int64 dense_rows = dense_shape(0);

    bool* empty_row_indicator = nullptr;
    if (context->output_required(kEmptyRowIndicatorOutput)) {
      Tensor* empty_row_indicator_t = nullptr;
      OP_REQUIRES_OK(context,
                     context->allocate_output(kEmptyRowIndicatorOutput,
                                              TensorShape({dense_rows}),
                                              &empty_row_indicator_t));
      empty_row_indicator = empty_row_indicator_t->vec<bool>().data();
    }
    int64* reverse_index_map = nullptr;
    if (context->output_required(kReverseIndexMapOutput)) {
      Tensor* reverse_index_map_t = nullptr;
      OP_REQUIRES_OK(context, context->allocate_output(kReverseIndexMapOutput,
                                                       TensorShape({N}),
                                                       &reverse_index_map_t));
      reverse_index_map = reverse_index_map_t->vec<int64>().data();
    }

    int rank = indices_t.shape().dim_size(1);

    if (dense_rows == 0) {
      OP_REQUIRES(
          context, N == 0,
          errors::InvalidArgument("Received SparseTensor with dense_shape[0] = "
                                  "0 but indices.shape[0] = ",
                                  N));
      Tensor* output_indices_t;
      TensorShape output_indices_shape({0, rank});
      OP_REQUIRES_OK(context, context->allocate_output(kOutputIndicesOutput,
                                                       output_indices_shape,
                                                       &output_indices_t));
      Tensor* output_values_t;
      OP_REQUIRES_OK(context, context->allocate_output(kOutputValuesOutput,
                                                       TensorShape({0}),
                                                       &output_values_t));

      // Exit early, nothing more to do.
      return;
    }

    bool rows_are_ordered = true;
    int64 last_indices_row = 0;
    std::vector<int64> csr_offset(dense_rows, 0);
    for (int i = 0; i < N; ++i) {
      const int64 row = indices(i, 0);
      OP_REQUIRES(context, row >= 0 && row < dense_rows,
                  errors::InvalidArgument("indices(", i, ", 0) is invalid: ",
                                          row, " >= ", dense_rows));
      ++csr_offset[row];
      rows_are_ordered = rows_are_ordered & (row >= last_indices_row);
      last_indices_row = row;
    }
    bool all_rows_full = true;
    for (int row = 0; row < dense_rows; ++row) {
      // csr_offset here describes the number of elements in this dense row
      bool row_empty = (csr_offset[row] == 0);
      if (empty_row_indicator) {
        empty_row_indicator[row] = row_empty;
      }
      all_rows_full = all_rows_full & !row_empty;
      // In filled version, each row has at least one element.
      csr_offset[row] = std::max(csr_offset[row], int64{1});
      // Update csr_offset to represent the number of elements up to and
      // including dense_row + 1:
      //  csr_offset(0) == #{elements of row 0}
      //  csr_offset(1) == #{elements of row 1} + #{elements of row 0}
      //  ..
      //  csr_offset(i) == starting index for elements in row i + 1.
      if (row > 0) {
        csr_offset[row] += csr_offset[row - 1];
      }
    }

    if (all_rows_full && rows_are_ordered) {
      context->set_output(kOutputIndicesOutput, indices_t);
      context->set_output(kOutputValuesOutput, values_t);
      if (reverse_index_map) {
        for (int64 i = 0; i < N; ++i) {
          reverse_index_map[i] = i;
        }
      }
    } else {
      Tensor* output_indices_t;
      const int64 N_full = csr_offset[dense_rows - 1];
      TensorShape output_indices_shape({N_full, rank});
      OP_REQUIRES_OK(context, context->allocate_output(kOutputIndicesOutput,
                                                       output_indices_shape,
                                                       &output_indices_t));
      auto output_indices = output_indices_t->matrix<int64>();

      Tensor* output_values_t;
      OP_REQUIRES_OK(context, context->allocate_output(kOutputValuesOutput,
                                                       TensorShape({N_full}),
                                                       &output_values_t));
      auto output_values = output_values_t->vec<T>();

      std::vector<int64> filled_count(dense_rows, 0);

      // Fill in values for rows that are not missing
      for (int64 i = 0; i < N; ++i) {
        const int64 row = indices(i, 0);
        int64& offset = filled_count[row];
        const int64 output_i = ((row == 0) ? 0 : csr_offset[row - 1]) + offset;
        offset++;  // Increment the filled count for this row.
        std::copy_n(&indices(i, 0), rank, &output_indices(output_i, 0));
        output_values(output_i) = values(i);
        // We'll need this reverse index map to backprop correctly.
        if (reverse_index_map) {
          reverse_index_map[i] = output_i;
        }
      }

      // Fill in values for rows that are missing
      for (int64 row = 0; row < dense_rows; ++row) {
        const int64 row_count = filled_count[row];
        if (row_count == 0) {  // We haven't filled this row
          const int64 starting_index = (row == 0) ? 0 : csr_offset[row - 1];
          // Remaining index values were set to zero already.
          // Just need to set the row index in the right location.
          output_indices(starting_index, 0) = row;
          for (int64 col = 1; col < rank; ++col) {
            output_indices(starting_index, col) = 0;
          }
          output_values(starting_index) = default_value;
        }
      }
    }
  }
};

#define REGISTER_KERNELS(type)                            \
  REGISTER_KERNEL_BUILDER(Name("SparseFillEmptyRows")     \
                              .Device(DEVICE_CPU)         \
                              .TypeConstraint<type>("T"), \
                          SparseFillEmptyRowsOp<type>)

TF_CALL_ALL_TYPES(REGISTER_KERNELS);
#undef REGISTER_KERNELS

template <typename T>
class SparseFillEmptyRowsGradOp : public OpKernel {
 public:
  explicit SparseFillEmptyRowsGradOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const Tensor* reverse_index_map_t;
    const Tensor* grad_values_t;
    OP_REQUIRES_OK(context,
                   context->input("reverse_index_map", &reverse_index_map_t));
    OP_REQUIRES_OK(context, context->input("grad_values", &grad_values_t));

    const CPUDevice& d = context->eigen_device<CPUDevice>();

    OP_REQUIRES(
        context, TensorShapeUtils::IsVector(reverse_index_map_t->shape()),
        errors::InvalidArgument("reverse_index_map must be a vector, saw: ",
                                reverse_index_map_t->shape().DebugString()));

    const auto reverse_index_map = reverse_index_map_t->vec<int64>();
    const auto grad_values = grad_values_t->vec<T>();

    const int64 N = reverse_index_map_t->shape().dim_size(0);
    const int64 N_full = grad_values_t->shape().dim_size(0);

    Tensor* d_values_t;
    OP_REQUIRES_OK(context, context->allocate_output(
                                "d_values", TensorShape({N}), &d_values_t));
    auto d_values = d_values_t->vec<T>();
    Tensor* d_default_value_t;
    OP_REQUIRES_OK(context,
                   context->allocate_output("d_default_value", TensorShape({}),
                                            &d_default_value_t));
    T& d_default_value = d_default_value_t->scalar<T>()();
    d_default_value = T();

    Tensor visited_t;
    OP_REQUIRES_OK(context, context->allocate_temp(
                                DT_BOOL, TensorShape({N_full}), &visited_t));
    auto visited = visited_t.vec<bool>();
    visited.device(d) = visited.constant(false);

    for (int i = 0; i < N; ++i) {
      // Locate the index of the output of the forward prop associated
      // with this location in the input of the forward prop.  Copy
      // the gradient into it.  Mark it as visited.
      d_values(i) = grad_values(reverse_index_map(i));
      visited(reverse_index_map(i)) = true;
    }
    for (int j = 0; j < N_full; ++j) {
      // The default value gradient gets the accumulated remainder of
      // the backprop values (since the default value was used to fill
      // in these slots in the forward calculation).
      if (!visited(j)) {
        d_default_value += grad_values(j);
      }
    }
  }
};

#define REGISTER_KERNELS(type)                            \
  REGISTER_KERNEL_BUILDER(Name("SparseFillEmptyRowsGrad") \
                              .Device(DEVICE_CPU)         \
                              .TypeConstraint<type>("T"), \
                          SparseFillEmptyRowsGradOp<type>)

TF_CALL_NUMBER_TYPES(REGISTER_KERNELS);
#undef REGISTER_KERNELS
}  // namespace tensorflow
<|endoftext|>


<commit_before>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#define EIGEN_USE_THREADS

#include <algorithm>
#include <numeric>
#include <unordered_map>
#include <utility>
#include <vector>

#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_util.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
#include "tensorflow/core/util/sparse/sparse_tensor.h"

namespace tensorflow {

using CPUDevice = Eigen::ThreadPoolDevice;

template <typename T>
class SparseFillEmptyRowsOp : public OpKernel {
 public:
  explicit SparseFillEmptyRowsOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const int kIndicesInput = 0;
    const int kValuesInput = 1;
    const int kDenseShapeInput = 2;
    const int kDefaultValueInput = 3;

    const int kOutputIndicesOutput = 0;
    const int kOutputValuesOutput = 1;
    const int kEmptyRowIndicatorOutput = 2;
    const int kReverseIndexMapOutput = 3;

    const Tensor& indices_t = context->input(kIndicesInput);
    const Tensor& values_t = context->input(kValuesInput);
    const Tensor& dense_shape_t = context->input(kDenseShapeInput);
    const Tensor& default_value_t = context->input(kDefaultValueInput);

    OP_REQUIRES(context, TensorShapeUtils::IsVector(dense_shape_t.shape()),
                errors::InvalidArgument("dense_shape must be a vector, saw: ",
                                        dense_shape_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsMatrix(indices_t.shape()),
                errors::InvalidArgument("indices must be a matrix, saw: ",
                                        indices_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsVector(values_t.shape()),
                errors::InvalidArgument("values must be a vector, saw: ",
                                        values_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(default_value_t.shape()),
                errors::InvalidArgument("default_value must be a scalar, saw: ",
                                        default_value_t.shape().DebugString()));
    // TODO(ebrevdo): add shape checks between values, indices,
    // dense_shape.  Also add check that dense rank > 0.

    const T& default_value = default_value_t.scalar<T>()();
    const auto indices = indices_t.matrix<int64>();
    const auto values = values_t.vec<T>();
    const auto dense_shape = dense_shape_t.vec<int64>();

    const int64 N = indices_t.shape().dim_size(0);
    const int64 dense_rows = dense_shape(0);

    bool* empty_row_indicator = nullptr;
    if (context->output_required(kEmptyRowIndicatorOutput)) {
      Tensor* empty_row_indicator_t = nullptr;
      OP_REQUIRES_OK(context,
                     context->allocate_output(kEmptyRowIndicatorOutput,
                                              TensorShape({dense_rows}),
                                              &empty_row_indicator_t));
      empty_row_indicator = empty_row_indicator_t->vec<bool>().data();
    }
    int64* reverse_index_map = nullptr;
    if (context->output_required(kReverseIndexMapOutput)) {
      Tensor* reverse_index_map_t = nullptr;
      OP_REQUIRES_OK(context, context->allocate_output(kReverseIndexMapOutput,
                                                       TensorShape({N}),
                                                       &reverse_index_map_t));
      reverse_index_map = reverse_index_map_t->vec<int64>().data();
    }

    int rank = indices_t.shape().dim_size(1);

    if (dense_rows == 0) {
      OP_REQUIRES(
          context, N == 0,
          errors::InvalidArgument("Received SparseTensor with dense_shape[0] = "
                                  "0 but indices.shape[0] = ",
                                  N));
      Tensor* output_indices_t;
      TensorShape output_indices_shape({0, rank});
      OP_REQUIRES_OK(context, context->allocate_output(kOutputIndicesOutput,
                                                       output_indices_shape,
                                                       &output_indices_t));
      Tensor* output_values_t;
      OP_REQUIRES_OK(context, context->allocate_output(kOutputValuesOutput,
                                                       TensorShape({0}),
                                                       &output_values_t));

      // Exit early, nothing more to do.
      return;
    }

    std::vector<int64> csr_offset(dense_rows, 0);
    for (int i = 0; i < N; ++i) {
      const int64 row = indices(i, 0);
      OP_REQUIRES(context, row >= 0 && row < dense_rows,
                  errors::InvalidArgument("indices(", i, ", 0) is invalid: ",
                                          row, " >= ", dense_rows));
      ++csr_offset[indices(i, 0)];
    }
    bool all_rows_full = true;
    for (int row = 0; row < dense_rows; ++row) {
      // csr_offset here describes the number of elements in this dense row
      bool row_empty = (csr_offset[row] == 0);
      if (empty_row_indicator) {
        empty_row_indicator[row] = row_empty;
      }
      all_rows_full = all_rows_full & !row_empty;
      // In filled version, each row has at least one element.
      csr_offset[row] = std::max(csr_offset[row], int64{1});
      // Update csr_offset to represent the number of elements up to and
      // including dense_row + 1:
      //  csr_offset(0) == #{elements of row 0}
      //  csr_offset(1) == #{elements of row 1} + #{elements of row 0}
      //  ..
      //  csr_offset(i) == starting index for elements in row i + 1.
      if (row > 0) {
        csr_offset[row] += csr_offset[row - 1];
      }
    }

    if (all_rows_full) {
      context->set_output(kOutputIndicesOutput, indices_t);
      context->set_output(kOutputValuesOutput, values_t);
      if (reverse_index_map) {
        for (int64 i = 0; i < N; ++i) {
          reverse_index_map[i] = i;
        }
      }
    } else {
      Tensor* output_indices_t;
      const int64 N_full = csr_offset[dense_rows - 1];
      TensorShape output_indices_shape({N_full, rank});
      OP_REQUIRES_OK(context, context->allocate_output(kOutputIndicesOutput,
                                                       output_indices_shape,
                                                       &output_indices_t));
      auto output_indices = output_indices_t->matrix<int64>();

      Tensor* output_values_t;
      OP_REQUIRES_OK(context, context->allocate_output(kOutputValuesOutput,
                                                       TensorShape({N_full}),
                                                       &output_values_t));
      auto output_values = output_values_t->vec<T>();

      std::vector<int64> filled_count(dense_rows, 0);

      // Fill in values for rows that are not missing
      for (int64 i = 0; i < N; ++i) {
        const int64 row = indices(i, 0);
        int64& offset = filled_count[row];
        const int64 output_i = ((row == 0) ? 0 : csr_offset[row - 1]) + offset;
        offset++;  // Increment the filled count for this row.
        std::copy_n(&indices(i, 0), rank, &output_indices(output_i, 0));
        output_values(output_i) = values(i);
        // We'll need this reverse index map to backprop correctly.
        if (reverse_index_map) {
          reverse_index_map[i] = output_i;
        }
      }

      // Fill in values for rows that are missing
      for (int64 row = 0; row < dense_rows; ++row) {
        const int64 row_count = filled_count[row];
        if (row_count == 0) {  // We haven't filled this row
          const int64 starting_index = (row == 0) ? 0 : csr_offset[row - 1];
          // Remaining index values were set to zero already.
          // Just need to set the row index in the right location.
          output_indices(starting_index, 0) = row;
          for (int64 col = 1; col < rank; ++col) {
            output_indices(starting_index, col) = 0;
          }
          output_values(starting_index) = default_value;
        }
      }
    }
  }
};

#define REGISTER_KERNELS(type)                            \
  REGISTER_KERNEL_BUILDER(Name("SparseFillEmptyRows")     \
                              .Device(DEVICE_CPU)         \
                              .TypeConstraint<type>("T"), \
                          SparseFillEmptyRowsOp<type>)

TF_CALL_ALL_TYPES(REGISTER_KERNELS);
#undef REGISTER_KERNELS

template <typename T>
class SparseFillEmptyRowsGradOp : public OpKernel {
 public:
  explicit SparseFillEmptyRowsGradOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const Tensor* reverse_index_map_t;
    const Tensor* grad_values_t;
    OP_REQUIRES_OK(context,
                   context->input("reverse_index_map", &reverse_index_map_t));
    OP_REQUIRES_OK(context, context->input("grad_values", &grad_values_t));

    const CPUDevice& d = context->eigen_device<CPUDevice>();

    OP_REQUIRES(
        context, TensorShapeUtils::IsVector(reverse_index_map_t->shape()),
        errors::InvalidArgument("reverse_index_map must be a vector, saw: ",
                                reverse_index_map_t->shape().DebugString()));

    const auto reverse_index_map = reverse_index_map_t->vec<int64>();
    const auto grad_values = grad_values_t->vec<T>();

    const int64 N = reverse_index_map_t->shape().dim_size(0);
    const int64 N_full = grad_values_t->shape().dim_size(0);

    Tensor* d_values_t;
    OP_REQUIRES_OK(context, context->allocate_output(
                                "d_values", TensorShape({N}), &d_values_t));
    auto d_values = d_values_t->vec<T>();
    Tensor* d_default_value_t;
    OP_REQUIRES_OK(context,
                   context->allocate_output("d_default_value", TensorShape({}),
                                            &d_default_value_t));
    T& d_default_value = d_default_value_t->scalar<T>()();
    d_default_value = T();

    Tensor visited_t;
    OP_REQUIRES_OK(context, context->allocate_temp(
                                DT_BOOL, TensorShape({N_full}), &visited_t));
    auto visited = visited_t.vec<bool>();
    visited.device(d) = visited.constant(false);

    for (int i = 0; i < N; ++i) {
      // Locate the index of the output of the forward prop associated
      // with this location in the input of the forward prop.  Copy
      // the gradient into it.  Mark it as visited.
      d_values(i) = grad_values(reverse_index_map(i));
      visited(reverse_index_map(i)) = true;
    }
    for (int j = 0; j < N_full; ++j) {
      // The default value gradient gets the accumulated remainder of
      // the backprop values (since the default value was used to fill
      // in these slots in the forward calculation).
      if (!visited(j)) {
        d_default_value += grad_values(j);
      }
    }
  }
};

#define REGISTER_KERNELS(type)                            \
  REGISTER_KERNEL_BUILDER(Name("SparseFillEmptyRowsGrad") \
                              .Device(DEVICE_CPU)         \
                              .TypeConstraint<type>("T"), \
                          SparseFillEmptyRowsGradOp<type>)

TF_CALL_NUMBER_TYPES(REGISTER_KERNELS);
#undef REGISTER_KERNELS
}  // namespace tensorflow
<commit_msg>Add judgment for whether the sp_input is ordered<commit_after>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#define EIGEN_USE_THREADS

#include <algorithm>
#include <numeric>
#include <unordered_map>
#include <utility>
#include <vector>

#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_util.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
#include "tensorflow/core/util/sparse/sparse_tensor.h"

namespace tensorflow {

using CPUDevice = Eigen::ThreadPoolDevice;

template <typename T>
class SparseFillEmptyRowsOp : public OpKernel {
 public:
  explicit SparseFillEmptyRowsOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const int kIndicesInput = 0;
    const int kValuesInput = 1;
    const int kDenseShapeInput = 2;
    const int kDefaultValueInput = 3;

    const int kOutputIndicesOutput = 0;
    const int kOutputValuesOutput = 1;
    const int kEmptyRowIndicatorOutput = 2;
    const int kReverseIndexMapOutput = 3;

    const Tensor& indices_t = context->input(kIndicesInput);
    const Tensor& values_t = context->input(kValuesInput);
    const Tensor& dense_shape_t = context->input(kDenseShapeInput);
    const Tensor& default_value_t = context->input(kDefaultValueInput);

    OP_REQUIRES(context, TensorShapeUtils::IsVector(dense_shape_t.shape()),
                errors::InvalidArgument("dense_shape must be a vector, saw: ",
                                        dense_shape_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsMatrix(indices_t.shape()),
                errors::InvalidArgument("indices must be a matrix, saw: ",
                                        indices_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsVector(values_t.shape()),
                errors::InvalidArgument("values must be a vector, saw: ",
                                        values_t.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(default_value_t.shape()),
                errors::InvalidArgument("default_value must be a scalar, saw: ",
                                        default_value_t.shape().DebugString()));
    // TODO(ebrevdo): add shape checks between values, indices,
    // dense_shape.  Also add check that dense rank > 0.

    const T& default_value = default_value_t.scalar<T>()();
    const auto indices = indices_t.matrix<int64>();
    const auto values = values_t.vec<T>();
    const auto dense_shape = dense_shape_t.vec<int64>();

    const int64 N = indices_t.shape().dim_size(0);
    const int64 dense_rows = dense_shape(0);

    bool* empty_row_indicator = nullptr;
    if (context->output_required(kEmptyRowIndicatorOutput)) {
      Tensor* empty_row_indicator_t = nullptr;
      OP_REQUIRES_OK(context,
                     context->allocate_output(kEmptyRowIndicatorOutput,
                                              TensorShape({dense_rows}),
                                              &empty_row_indicator_t));
      empty_row_indicator = empty_row_indicator_t->vec<bool>().data();
    }
    int64* reverse_index_map = nullptr;
    if (context->output_required(kReverseIndexMapOutput)) {
      Tensor* reverse_index_map_t = nullptr;
      OP_REQUIRES_OK(context, context->allocate_output(kReverseIndexMapOutput,
                                                       TensorShape({N}),
                                                       &reverse_index_map_t));
      reverse_index_map = reverse_index_map_t->vec<int64>().data();
    }

    int rank = indices_t.shape().dim_size(1);

    if (dense_rows == 0) {
      OP_REQUIRES(
          context, N == 0,
          errors::InvalidArgument("Received SparseTensor with dense_shape[0] = "
                                  "0 but indices.shape[0] = ",
                                  N));
      Tensor* output_indices_t;
      TensorShape output_indices_shape({0, rank});
      OP_REQUIRES_OK(context, context->allocate_output(kOutputIndicesOutput,
                                                       output_indices_shape,
                                                       &output_indices_t));
      Tensor* output_values_t;
      OP_REQUIRES_OK(context, context->allocate_output(kOutputValuesOutput,
                                                       TensorShape({0}),
                                                       &output_values_t));

      // Exit early, nothing more to do.
      return;
    }

    bool indices_is_order = true;
    int64 last_indices_row = 0;
    std::vector<int64> csr_offset(dense_rows, 0);
    for (int i = 0; i < N; ++i) {
      const int64 row = indices(i, 0);
      OP_REQUIRES(context, row >= 0 && row < dense_rows,
                  errors::InvalidArgument("indices(", i, ", 0) is invalid: ",
                                          row, " >= ", dense_rows));
      ++csr_offset[row];
      indices_is_order = indices_is_order & (row >= last_indices_row);
      last_indices_row = row;
    }
    bool all_rows_full = true;
    for (int row = 0; row < dense_rows; ++row) {
      // csr_offset here describes the number of elements in this dense row
      bool row_empty = (csr_offset[row] == 0);
      if (empty_row_indicator) {
        empty_row_indicator[row] = row_empty;
      }
      all_rows_full = all_rows_full & !row_empty;
      // In filled version, each row has at least one element.
      csr_offset[row] = std::max(csr_offset[row], int64{1});
      // Update csr_offset to represent the number of elements up to and
      // including dense_row + 1:
      //  csr_offset(0) == #{elements of row 0}
      //  csr_offset(1) == #{elements of row 1} + #{elements of row 0}
      //  ..
      //  csr_offset(i) == starting index for elements in row i + 1.
      if (row > 0) {
        csr_offset[row] += csr_offset[row - 1];
      }
    }

    if (all_rows_full && indices_is_order) {
      context->set_output(kOutputIndicesOutput, indices_t);
      context->set_output(kOutputValuesOutput, values_t);
      if (reverse_index_map) {
        for (int64 i = 0; i < N; ++i) {
          reverse_index_map[i] = i;
        }
      }
    } else {
      Tensor* output_indices_t;
      const int64 N_full = csr_offset[dense_rows - 1];
      TensorShape output_indices_shape({N_full, rank});
      OP_REQUIRES_OK(context, context->allocate_output(kOutputIndicesOutput,
                                                       output_indices_shape,
                                                       &output_indices_t));
      auto output_indices = output_indices_t->matrix<int64>();

      Tensor* output_values_t;
      OP_REQUIRES_OK(context, context->allocate_output(kOutputValuesOutput,
                                                       TensorShape({N_full}),
                                                       &output_values_t));
      auto output_values = output_values_t->vec<T>();

      std::vector<int64> filled_count(dense_rows, 0);

      // Fill in values for rows that are not missing
      for (int64 i = 0; i < N; ++i) {
        const int64 row = indices(i, 0);
        int64& offset = filled_count[row];
        const int64 output_i = ((row == 0) ? 0 : csr_offset[row - 1]) + offset;
        offset++;  // Increment the filled count for this row.
        std::copy_n(&indices(i, 0), rank, &output_indices(output_i, 0));
        output_values(output_i) = values(i);
        // We'll need this reverse index map to backprop correctly.
        if (reverse_index_map) {
          reverse_index_map[i] = output_i;
        }
      }

      // Fill in values for rows that are missing
      for (int64 row = 0; row < dense_rows; ++row) {
        const int64 row_count = filled_count[row];
        if (row_count == 0) {  // We haven't filled this row
          const int64 starting_index = (row == 0) ? 0 : csr_offset[row - 1];
          // Remaining index values were set to zero already.
          // Just need to set the row index in the right location.
          output_indices(starting_index, 0) = row;
          for (int64 col = 1; col < rank; ++col) {
            output_indices(starting_index, col) = 0;
          }
          output_values(starting_index) = default_value;
        }
      }
    }
  }
};

#define REGISTER_KERNELS(type)                            \
  REGISTER_KERNEL_BUILDER(Name("SparseFillEmptyRows")     \
                              .Device(DEVICE_CPU)         \
                              .TypeConstraint<type>("T"), \
                          SparseFillEmptyRowsOp<type>)

TF_CALL_ALL_TYPES(REGISTER_KERNELS);
#undef REGISTER_KERNELS

template <typename T>
class SparseFillEmptyRowsGradOp : public OpKernel {
 public:
  explicit SparseFillEmptyRowsGradOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const Tensor* reverse_index_map_t;
    const Tensor* grad_values_t;
    OP_REQUIRES_OK(context,
                   context->input("reverse_index_map", &reverse_index_map_t));
    OP_REQUIRES_OK(context, context->input("grad_values", &grad_values_t));

    const CPUDevice& d = context->eigen_device<CPUDevice>();

    OP_REQUIRES(
        context, TensorShapeUtils::IsVector(reverse_index_map_t->shape()),
        errors::InvalidArgument("reverse_index_map must be a vector, saw: ",
                                reverse_index_map_t->shape().DebugString()));

    const auto reverse_index_map = reverse_index_map_t->vec<int64>();
    const auto grad_values = grad_values_t->vec<T>();

    const int64 N = reverse_index_map_t->shape().dim_size(0);
    const int64 N_full = grad_values_t->shape().dim_size(0);

    Tensor* d_values_t;
    OP_REQUIRES_OK(context, context->allocate_output(
                                "d_values", TensorShape({N}), &d_values_t));
    auto d_values = d_values_t->vec<T>();
    Tensor* d_default_value_t;
    OP_REQUIRES_OK(context,
                   context->allocate_output("d_default_value", TensorShape({}),
                                            &d_default_value_t));
    T& d_default_value = d_default_value_t->scalar<T>()();
    d_default_value = T();

    Tensor visited_t;
    OP_REQUIRES_OK(context, context->allocate_temp(
                                DT_BOOL, TensorShape({N_full}), &visited_t));
    auto visited = visited_t.vec<bool>();
    visited.device(d) = visited.constant(false);

    for (int i = 0; i < N; ++i) {
      // Locate the index of the output of the forward prop associated
      // with this location in the input of the forward prop.  Copy
      // the gradient into it.  Mark it as visited.
      d_values(i) = grad_values(reverse_index_map(i));
      visited(reverse_index_map(i)) = true;
    }
    for (int j = 0; j < N_full; ++j) {
      // The default value gradient gets the accumulated remainder of
      // the backprop values (since the default value was used to fill
      // in these slots in the forward calculation).
      if (!visited(j)) {
        d_default_value += grad_values(j);
      }
    }
  }
};

#define REGISTER_KERNELS(type)                            \
  REGISTER_KERNEL_BUILDER(Name("SparseFillEmptyRowsGrad") \
                              .Device(DEVICE_CPU)         \
                              .TypeConstraint<type>("T"), \
                          SparseFillEmptyRowsGradOp<type>)

TF_CALL_NUMBER_TYPES(REGISTER_KERNELS);
#undef REGISTER_KERNELS
}  // namespace tensorflow
<|endoftext|>