Checkpoint and Function Composition: Example and Test

checkpoint.cpp

Checkpoint and Function Composition: Example and Test
Purpose
Break a large computation into pieces and only store values at the interface of the pieces. In actual applications, there may be many functions, but for this example there are only two. The functions

F : R^{2} \to R^{2}

and

G : R^{2} \to R^{2}

defined by

F (x) = (\begin{matrix} x_{0} x_{1} \\ x_{1} - x_{0} \end{matrix}), G (y) = (\begin{matrix} y_{0} - y_{1} \\ y_{1} y_{0} \end{matrix})

Processing Steps
We apply reverse mode to compute the derivative of

H : R^{2} \to R

is defined by

\begin{matrix} H (x) & = & G_{0} [F (x)] + G_{1} [F (x)] \\ = & x_{0} x_{1} - (x_{1} - x_{0}) + x_{0} x_{1} (x_{1} - x_{0}) \\ = & x_{0} x_{1} (1 - x_{0} + x_{1}) - x_{1} + x_{0} \end{matrix}

Given the zero and first order Taylor coefficients

x^{(0)}

and

x^{(1)}

, we use

X (t)

Y (t)

and

Z (t)

for the corresponding functions; i.e.,

\begin{matrix} X (t) & = & x^{(0)} + x^{(1)} t \\ Y (t) & = & F [X (t)] = y^{(0)} + y^{(1)} t + O (t^{2}) \\ Z (t) & = & G {F [X (t)]} = z^{(0)} + z^{(1)} t + O (t^{2}) \\ h^{(0)} & = & z_{0}^{(0)} + z_{1}^{(0)} \\ h^{(1)} & = & z_{0}^{(1)} + z_{1}^{(1)} \end{matrix}

Here are the processing steps:

Use forward mode on $F (x)$ to compute $y^{(0)}$ and $y^{(1)}$ .
Free some, or all, of the memory corresponding to $F (x)$ .
Use forward mode on $G (y)$ to compute $z^{(0)}$ and $z^{(1)}$
Use reverse mode on $G (y)$ to compute the derivative of $h^{(1)}$ with respect to $y^{(0)}$ and $y^{(1)}$ .
Free all the memory corresponding to $G (y)$ .
Use reverse mode on $F (x)$ to compute the derivative of $h^{(1)}$ with respect to $x^{(0)}$ and $x^{(1)}$ .

This uses the following relations:

\begin{matrix} \partial_{x (0)} h^{(1)} [x^{(0)}, x^{(1)}] & = & \partial_{y (0)} h^{(1)} [y^{(0)}, y^{(1)}] \partial_{x (0)} y^{(0)} [x^{(0)}, x^{(1)}] \\ + & \partial_{y (1)} h^{(1)} [y^{(0)}, y^{(1)}] \partial_{x (0)} y^{(1)} [x^{(0)}, x^{(1)}] \\ \partial_{x (1)} h^{(1)} [x^{(0)}, x^{(1)}] & = & \partial_{y (0)} h^{(1)} [y^{(0)}, y^{(1)}] \partial_{x (1)} y^{(0)} [x^{(0)}, x^{(1)}] \\ + & \partial_{y (1)} h^{(1)} [y^{(0)}, y^{(1)}] \partial_{x (1)} y^{(1)} [x^{(0)}, x^{(1)}] \end{matrix}

where

\partial_{x (0)}

denotes the partial with respect to

x^{(0)}


 

# include <cppad/cppad.hpp>

namespace {
	template <class Vector>
	Vector F(const Vector& x)
	{	Vector y(2);
		y[0] = x[0] * x[1];
		y[1] = x[1] - x[0];
		return y;
	}
	template <class Vector>
	Vector G(const Vector& y)
	{	Vector z(2);
		z[0] = y[0] - y[1];
		z[1] = y[1] * y[0];
		return z;
	}
}

namespace {
	bool checkpoint_case(bool free_all)
	{	bool ok = true;
     	double eps = 10. * CppAD::epsilon<double>();

		using CppAD::AD;
		using CppAD::NearEqual;
		CppAD::ADFun<double> f, g, empty;

		// specify the Taylor coefficients for X(t)
		size_t n    = 2;
		CPPAD_TEST_VECTOR<double> x0(n), x1(n);
		x0[0] = 1.; x0[1] = 2.;
		x1[0] = 3.; x1[1] = 4.;

		// record the function F(x)
		CPPAD_TEST_VECTOR< AD<double> > X(n), Y(n);
		size_t i;
		for(i = 0; i < n; i++)
			X[i] = x0[i];
		CppAD::Independent(X);
		Y = F(X);
		f.Dependent(X, Y);

		// a fucntion object with an almost empty operation sequence
		CppAD::Independent(X);
		empty.Dependent(X, X);

		// compute the Taylor coefficients for Y(t)
		CPPAD_TEST_VECTOR<double> y0(n), y1(n);
		y0 = f.Forward(0, x0);
		y1 = f.Forward(1, x1);
		if( free_all )
			f = empty;
		else	
		{	// free all the Taylor coefficients stored in f
			f.capacity_taylor(0);
		}

		// record the function G(x)
		CPPAD_TEST_VECTOR< AD<double> > Z(n);
		CppAD::Independent(Y);
		Z = G(Y);
		g.Dependent(Y, Z);

		// compute the Taylor coefficients for Z(t)
		CPPAD_TEST_VECTOR<double> z0(n), z1(n);
		z0 = g.Forward(0, y0);
		z1 = g.Forward(1, y1);

		// check zero order Taylor coefficient for h^0 = z_0^0 + z_1^0
		double check = x0[0] * x0[1] * (1. - x0[0] + x0[1]) - x0[1] + x0[0];
		double h0    = z0[0] + z0[1];
		ok          &= NearEqual(h0, check, eps, eps);

		// check first order Taylor coefficient h^1
		check     = x0[0] * x0[1] * (- x1[0] + x1[1]) - x1[1] + x1[0];
		check    += x1[0] * x0[1] * (1. - x0[0] + x0[1]);
		check    += x0[0] * x1[1] * (1. - x0[0] + x0[1]);
		double h1 = z1[0] + z1[1];
		ok       &= NearEqual(h1, check, eps, eps);
	
		// compute the derivative with respect to y^0 and y^0 of h^1
		size_t p = 2;
		CPPAD_TEST_VECTOR<double> w(n*p), dw(n*p);
		w[0*p+0] = 0.; // coefficient for z_0^0
		w[0*p+1] = 1.; // coefficient for z_0^1
		w[1*p+0] = 0.; // coefficient for z_1^0
		w[1*p+1] = 1.; // coefficient for z_1^1
		dw       = g.Reverse(p, w);

		// We are done using g, so we can free its memory.
		g = empty;
		// We need to use f next. 
		if( free_all )
		{	// we must again record the operation sequence for F(x)
			CppAD::Independent(X);
			Y = F(X);
			f.Dependent(X, Y);
		}
		// now recompute the Taylor coefficients corresponding to F(x)
		// (we already know the result; i.e., y0 and y1).
		f.Forward(0, x0);
		f.Forward(1, x1);

		// compute the derivative with respect to x^0 and x^0 of
		//	h^1 = z_0^1 + z_1^1
		CPPAD_TEST_VECTOR<double> dv(n*p);
		dv   = f.Reverse(p, dw); 

		// check partial of h^1 w.r.t x^0_0
		check  = x0[1] * (- x1[0] + x1[1]);
		check -= x1[0] * x0[1];
		check += x1[1] * (1. - x0[0] + x0[1]) - x0[0] * x1[1];
		ok    &= NearEqual(dv[0*p+0], check, eps, eps);

		// check partial of h^1 w.r.t x^0_1
		check  = x0[0] * (- x1[0] + x1[1]);
		check += x1[0] * (1. - x0[0] + x0[1]) + x1[0] * x0[1]; 
		check += x0[0] * x1[1];
		ok    &= NearEqual(dv[1*p+0], check, eps, eps);

		// check partial of h^1 w.r.t x^1_0
		check  = 1. - x0[0] * x0[1];
		check += x0[1] * (1. - x0[0] + x0[1]);
		ok    &= NearEqual(dv[0*p+1], check, eps, eps);

		// check partial of h^1 w.r.t x^1_1
		check  = x0[0] * x0[1] - 1.;
		check += x0[0] * (1. - x0[0] + x0[1]);
		ok    &= NearEqual(dv[1*p+1], check, eps, eps);

		return ok;
	}
}
bool checkpoint(void)
{	bool ok = true;
	ok     &= checkpoint_case(true);
	ok     &= checkpoint_case(false);
	return ok;
}

Input File: example/checkpoint.cpp