Fast Automatic Differentiation in C++

//////////////////////////////////////////////////////////////////////////
/// The Reed Library 
/// Copyright (c) 2006 Yevgeniy Shatokhin 

/// File: asection.h 
/// Contains definitions necessary for active section manipulation.

/// 
/// Last update: August 29, 2006

#ifndef __ASECTION_H__E079034E_E787_4946_8B13_17202FD9FEAD_INCLUDED_
#define __ASECTION_H__E079034E_E787_4946_8B13_17202FD9FEAD_INCLUDED_

/// Define REED_MULTITHREADED if you intend to use the Reed-instrumented
/// code in several threads sumultaneously. Otherwise it's better to 
/// undefine this macro to (probably) get better performance.
#ifdef REED_MULTITHREADED
#define REED_THREAD __declspec(thread)
#else
#define REED_THREAD
#endif	// #ifdef REED_MULTITHREADED

// I do not like bools, so I prefer to use ints instead.
// If TRUE and FALSE are not defined we define them here.
#ifndef FALSE
#define FALSE 0
#define TRUE !FALSE
#endif

#include <reed/mem_pool.h>	// special allocator
#include <reed/node.h>		// CG node classes
#include <cassert>
#include <vector>

namespace reed
{
	class CADouble;	// forward declaration

	//////////////////////////////////////////////////////////////////////////
	/// CActiveSection class \n
	/// Instances of this class represent an "active section" (AS), a program
	/// fragment the structure of which is to be recorded in the CG (computational
	/// graph, "graf algoritma" in russian).
	//////////////////////////////////////////////////////////////////////////
	class CActiveSection
	{
	public:
		/// 
		/// Number of code execution modes. (See 'EMode' below for their descriptions.)
		enum {NMODES = 4};

		///
		/// Type for the current mode of code execution in this thread.
		enum EMode
		{
			/// executing code outside any AS (this mode cannot be assigned to
			/// any particular AS, it's just for convenience).
			AS_MODE_OUTSIDE		= 0,	
			AS_MODE_PLAIN		= 1,	///< executing code within this AS, recording new CG
			AS_MODE_CACHING1	= 2,	///< executing code within this AS, using cached CG
			AS_MODE_CACHING2	= 3		///< (reserved)
		};

	public:
		// Static methods

		/// 
		/// Get current execution mode.
		static EMode getMode()
		{ 
			return (activeAS == NULL) ? AS_MODE_OUTSIDE : getCurrentAS().mode; 
		}

		///
		/// Get a reference to the current AS (for the current thread).
		static CActiveSection& getCurrentAS()
		{ 
			assert(activeAS != NULL);
			return (*activeAS);
		}

	public:
		// typedefs
		typedef std::vector<CADouble> TADVector;
		typedef std::vector<double> TDoubleVector;

	public:
		// Non-static service methods.

		///
		/// Constructor.
		CActiveSection();

		/// 
		/// Copy constructor. Not sure if it's really necessary. 
        CActiveSection(const CActiveSection& rhs);

		///
		/// Destructor.
		~CActiveSection();

		///
		/// Return TRUE if CG caching is enabled, FALSE otherwise.
		int isCacheEnabled() const
		{ return bCacheEnabled; }

		///
		/// Set AS recording mode to plain (no caching)
		void setPlainMode()
		{
			mode = AS_MODE_PLAIN;
		}

		///
		/// Pass TRUE to enable or FALSE to disable CG caching. \n
		/// WARNING. This method can be called from outside this AS only.
		void enableCache(int bEnable = TRUE)
		{ 
			assert(activeAS != this);	// We must not be within this AS.
			if (bEnable) 
			{
				bCacheEnabled = TRUE;
			}
			else
			{
				bCacheEnabled = FALSE;
				mode = AS_MODE_PLAIN;
			}
		}

		/// 
		/// Mark the beginning of the AS.
		void begin();

		///
		/// Mark the end of the AS.
		void end();

		/// Add a copy of the given node to the CG. Pointers to the arguments
		/// should be already set.
		template <CNodeBase::EOpCode _Op>
		CNodeBase* addNode(const CNode<_Op>& node)
		{
			CNode<_Op>* p = static_cast<CNode<_Op>*>(alloc.allocate());

			*((unsigned long*)p) = *((unsigned long*)&node); //__vfptr, always the 1st 4 bytes
			p->left = node.left;
			p->right = node.right;
			p->value.d = node.value.d;
			
			return p;
		}

		/// Retrieve accumulated values from the CG nodes associated with 
		/// the variables listed in the arg array.\n
		/// The previous contents of the 'val' array will be deleted in this function.
		void getAccumulators(TADVector& arg, TDoubleVector& val);

		/// Set accumulated values to the CG nodes associated with the variables 
		/// listed in the arg array.\n
		/// The 'val' vector contains the new values to set. 'arg' and 'val' must
		/// have equal number of the elements.
		void setAccumulators(TADVector& arg, TDoubleVector& val);

	public:
		// algorithms operating on the AS

		/// Update (if possible) the recorded CG of this active section.\n
        /// Arguments:\n
		/// [in] arg - a vector of the independent ('input') variables that 
		/// will have new values.\n
		/// [in] val - new values for these variables (in the same order).\n
		///	The independent variables not included in ths vector retain their
		///	previous values.\n
		/// 'arg' and 'val' vectors must have the same size.\n
		/// Return value - TRUE if the update was successful (the CG structure
		/// is the same for new values as for the old ones), FALSE if any of the
		/// relational operators changed its value. \n
		/// If update() succeeds, 'value' fields in the CG nodes are updated. 
		/// Otherwise the function has no effect, 'value' fields remain the same.\n
		/// (NB) This doesn't concern the accumulators. Their contents may (and most
		/// likely will) change in course of this function.
		/// You can use getAccumulators() to save the accumulated values of interest
		/// to a different location before calling update.
		int update(TADVector& arg, TDoubleVector& val);

		/// Perform the 1st order scalar FAD (fast automatic differentiation) in the
		/// forward mode.\n
		/// Arguments:\n
		/// [in] arg - a vector of the independent ('input') variables of interest.
		/// The differentiation will take place against these.\n
		/// [in] s - a vector of multipliers for these variables (in the same order) - 
		///	see below.\n
		/// 'arg' and 's' vectors must have the same size.\n
		/// Return value: status of derivative computation (see CNodeBase::EDerivStatus).
		/// After this 'forward pass' each node's accumulator contains the inner (dot)
		/// product of s vector and the gradient of the function F, which is defined 
		/// by a subgraph of the CG beginning from the nodes associated with 
		/// arg[0], arg[1], ... etc. and ending with this particular node.\n
		/// Example. \n
		/// Let y(u, v) be defined as follows: y = f(g(u, v)). arg[0]=u, arg[1]=v.
		/// After the forward pass the nodes will contain the following in their accumulators:
		/// (except the nodes for relational operators)\n
		/// acc(u) = s[0], acc(v) = s[1], acc(g) = s[0]*(dg/du)+s[1]*(dg/dv),
		/// acc(f) = s[0]*(df/dg)*(dg/du)+s[1]*(df/dg)*(dg/dv), \n
		/// i.e. acc(f)=s[0]*(dy/du)+s[1]*(dy/dv).\n 
		/// If w[i]==1, s[j]==0, for all j!=i then acc(f) = dy/darg[i].\n
		/// The forward mode allows to compute derivatives of all functions of interest
		/// against one independent variable in a single pass.
		int forward(TADVector& arg, TDoubleVector& s);

		/// Compute a derivative of the 'func' function against the independent variable 
		/// 'arg' and return it in 'out'. If arg is not independent the result is undefined.
		/// Return value - see 'forward(TADVector& arg, TDoubleVector& s)'.
		int forward(CADouble arg, CADouble func, double& out);

		/// Compute the inner products of 's' and gradients of the functions from 'func'
		/// against the independent varibles from 'arg'. The results are stored in 'out'.\n
		/// For further description of the return value, s and arg see 
		/// 'forward(TADVector& arg, TDoubleVector& s)'.\n
		/// Ex. out[0] = s[0]*(dfunc[0]/darg[0]) + s[1]*(dfunc[0]/darg[1]) + ...\n
		/// (NB) Previous contents of 'out' vector will be cleared in this function.
		int forward(TADVector& arg, TADVector& func, TDoubleVector& s, TDoubleVector& out);

		/// Perform the 1st order scalar FAD (fast automatic differentiation) in the
		/// reverse mode.\n
		/// The CG must contain at least one node.\n
		/// Arguments:\n
		/// [in] func - a vector of the dependent ('output') variables of interest.
		/// It's the derivatives of these we want to get.\n
		/// [in] w - a weight vector for these functions (in the same order) - see below.\n
		/// 'func' and 'w' vectors must have the same size.\n
		/// Return value: status of derivative computation (see CNodeBase::EDerivStatus).
		/// After this 'reverse pass' each node (x) contains in the accumulator the inner (dot)
		/// product of w and (dfunc[0]/dx, dfunc[1]/dx, ...) 
		/// (except the nodes for relational operators)\n
		/// Example. \n
		/// Let y(u, v) and z(u, v) be defined as follows: 
		/// y = f(g(u, v)), z = h(g(u, v)); func[0]=f, func[1]=h.
		/// After the reverse pass the nodes will contain the following in their accumulators:\n
		/// acc(f)=w[0], acc(h)=w[1], acc(g)=w[0]*(df/dg)+w[1]*(dh/dg),
		/// acc(u)=w[0]*(df/dg)*(dg/du)+w[1]*(dh/dg)*(dg/du), acc(v)=w[0]*(df/dg)*(dg/dv)+w[1]*(dh/dg)*(dg/dv).
		/// i.e. 
		/// If w[0]==1, w[1]==0, then acc(u) = dy/du, acc(v) = dy/dv.
		/// The reverse mode allows to compute the complete gradient of the given function
		/// in a single pass.
		int reverse(TADVector& func, TDoubleVector& w);

		/// Compute a derivative of the 'func' function against the variable 
		/// 'arg' and return it in 'out'. If arg can be a dependent variable.
		/// Return value - see 'reverse(TADVector& func, TDoubleVector& w)'.
		int reverse(CADouble arg, CADouble func, double& out);

		/// Compute the 'weighted' derivatives of the functions from 'func'
		/// against the varibles from 'arg'. The results are stored in 'out'.\n
		/// For further description of the return value, w and arg see 
		/// 'reverse(TADVector& func, TDoubleVector& w)'.
		/// Ex. out[0] = w[0]*(dfunc[0]/darg[0]) + w[1]*(dfunc[1]/darg[0]) + ...\n
		/// (NB) Previous contents of 'out' vector will be cleared in this function.
		int reverse(TADVector& arg, TADVector& func, TDoubleVector& w, TDoubleVector& out);

		/// Compute lower estimate of rounding error using a 'forward-like' pass.\n
		/// Arguments:\n
		/// [in] arg - a vector of the independent ('input') variables of interest,
		/// for which we'd like to specify initial errors different from the default.
		/// [in] err - a vector of these initial errors for the variables from 'arg' 
		/// (in the same order). 'arg' and 'err' must have the same size.\n
		/// If an independent variable is not mentioned in 'arg', default initial error
		/// is used for this variable. The default error is the 'storage error', derr(x).
		/// See reed::derr(x) in util.h for more information. 
		/// Pass empty arg and err if you don't want to overwrite default error values.\n
		/// Return value: status of derivative computation (see CNodeBase::EDerivStatus).\n
		/// When the function returns, each node (except those for relational operators)
		/// contains in its accumulator the lower error estimate of the value in this node.\n
		/// It's assumed there are 2 main sources of rounding errors:\n
		/// 1. 'Storage error' mentioned above that is due to a limited width of the double's
		/// bit representation. (According to the IEEE standard, of all 64 bits in a double
		/// one is used to store a sign, 11 bits hold the order and 52 are left for mantissa).\n
		/// 2. 'Uncertainty error' that arises from any math operation on inexact arguments,
		/// that is, the arguments that possibly have values only known with some error.\n
		/// E.g., if a == 1 (exactly) than sin(a) == sin(1) exactly. However, if all we know
		/// about a is that it resides in some interval like [1.0 - da, 1.0 + da],  
		/// the sin(a) value can be only estimated rather than computed exactly.\n
		/// It's also assumed that the internal computation errors are negligible (i.e. less
		/// than storage errors). That means the number of trusted bits in the IEEE-representation
		/// of x and sin(x) is the same. Not quite sure if this is the case, but let's assume
		/// that for now at least.\n
        int ree_lower(TADVector& arg, TDoubleVector& err);

		/// Compute the lower error estimates for the functions from 'func'.
		/// The independent varibles from 'arg' get their initial error estimates equal
		/// to those specified in 'err'. The results are stored in 'out'.\n
		/// For further description of the return value, err and arg see 
		/// 'ree_lower(TADVector& arg, TDoubleVector& err)'.\n
		/// (NB) Previous contents of 'out' vector will be cleared in this function.
		int ree_lower(TADVector& arg, TADVector& func, TDoubleVector& err, TDoubleVector& out);

	private:
		///
		/// The active section we are in (NULL if there is none).
		REED_THREAD static CActiveSection* activeAS;

	private:
		///
		/// Current mode of code execution regarding this AS (the "state" of the AS).
		EMode mode;

		///
		/// TRUE if automatic caching is allowed for the CG in this AS.
		int bCacheEnabled;

	public:
		CMemoryPool<CNodeBase> alloc; ///< Special allocator for CG nodes.

		CNodeBase* pCurrent; ///< Pointer to current node (used to implement caching).
	};

} // namespace reed

#endif //__ASECTION_H__E079034E_E787_4946_8B13_17202FD9FEAD_INCLUDED_