Print["This is Perturbation AIM Full Version 2.4.1ets"] ;
(*                                                                          *)
(* Credits: original algorithm designed by Gary Anderson, 2002-3;           *)
(*   corrected, optimized, and recoded by Eric Swanson, 2004.               *)
(* Code is currently maintained by Eric Swanson, see                        *)
(*   http://www.ericswanson.pro for most recent version and for             *)
(*   background, instructions, capabilities, examples, and tips.            *)
(* This code will only work with Mathematica version 5.  While it is not    *)
(*   too hard to modify the code to work in earlier versions of             *)
(*   Mathematica, those earlier versions are *much* slower at solving       *)
(*   linear systems of equations, and a few of the function calls are also  *)
(*   a bit more awkward.                                                    *)
(* This is the "advanced" or "Full" version of the perturbationAIM code.    *)
(*   It is generally faster and more immune to potential numerical          *)
(*   problems than the "simple" version, but it is also longer and less     *)
(*   transparent.  If you want to understand the underlying algorithm       *)
(*   and the implementation of the algorithm, you will find it much         *)
(*   easier to work through the "simple" version of the code,               *)
(*   perturbationAIMsimple.m, first, and then go through the code below.    *)
(*   The "simple" version is available for download from the web site       *)
(*   above.                                                                 *)
(* Both the "Full" and "simple" versions should always yield identical      *)
(*   solutions to a given model, up to issues of machine precision and      *)
(*   numerical accuracy.  The "simple" version should be completely         *)
(*   adequate for most purposes, however, and will be much easier for the   *)
(*   user to understand and even tailor to his/her individual needs.        *)
(* No matter which version you are using, you should compute your final     *)
(*   set of results using the arbitrary precision capabilities of the       *)
(*   code, as these are *much* more likely to be free of numerical          *)
(*   problems.                                                              *)
(*                                                                          *)

SetOptions["stdout",PageWidth->79] ;
<<DiscreteMath`Combinatorica`
<<LinearAlgebra`MatrixManipulation`

AIMZeroTol=10^-10 ;


(* Useful utilities that we will repeatedly call.                           *)
(* The primary improvement here with respect to the "simple" version is     *)
(*   the introduction of a new routine, AIMBackLookingVars, that searches   *)
(*   through the model for any variables that are clearly "backward-        *)
(*   looking" (i.e., obviously a function of just a few state variables,    *)
(*   typically one or two, rather than the whole set of state variables     *)
(*   of the model).  You could call these clearly "predetermined"           *)
(*   variables if you like, but I've avoided using that term.  This helps   *)
(*   speed up the code by reducing the amount of differentiation that       *)
(*   must be done and imposing a number of zero restrictions on the         *)
(*   solution.                                                              *)

AIMGenericArgs[eqns_]:= AIMGenericArgs[eqns] =
Table[Unique["arg"],{AIMNArgs[eqns]}] ;

AIMLagVars[eqns_]:= AIMLagVars[eqns] =
Flatten[Map[Table[#[t+i],{i,Min[Cases[eqns,#[t+j_.]->j,Infinity]],-1}]&,
							AIMVarNames[eqns]]] ;
AIMMaxLag[eqns_]:= AIMMaxLag[eqns] =
-Min[Cases[eqns,x_Symbol[t+i_.]->i,Infinity]] ;

AIMMaxLead[eqns_]:= AIMMaxLead[eqns] =
Max[Cases[eqns,x_Symbol[t+i_.]->i,Infinity]] ;

AIMNArgs[eqns_]:= AIMNArgs[eqns] = Length[AIMStateVars[eqns]] ;

AIMNEqns[eqns_]:= AIMNEqns[eqns] = Length[Flatten[{eqns}]] ;

AIMNVars[eqns_]:= AIMNVars[eqns] = Length[AIMVarNames[eqns]] ;

AIMShocks[eqns_]:= AIMShocks[eqns] = Union[Cases[eqns,eps[_][t],Infinity]] ;

AIMSSSubs={Sigma->0, eps[_][_]->0, x_[t+_.]:>Symbol[SymbolName[x]<>"AIMSS"]} ;

AIMSSVars[eqns_]:= AIMSSVars[eqns] = 
Map[Symbol[SymbolName[#]<>"AIMSS"]&, AIMVarNames[eqns]] ;

AIMStateVars[eqns_]:= AIMStateVars[eqns] =
Join[AIMLagVars[eqns], AIMShocks[eqns], {Sigma}] ;

AIMVarNames[eqns_]:= AIMVarNames[eqns] =
Union[Cases[eqns,x_Symbol[t+i_.]->x,Infinity]] ;

(* Here is the AIMBackLookingVars routine.  It first finds all equations    *)
(*   that have no leads.  For those equations that have only one variable   *)
(*   dated t, it is then obvious that the solution to this variable is      *)
(*   entirely a function of the lagged variables and shocks in that         *)
(*   particular equation.  This particular variable can then be solved      *)
(*   in terms of a reduced state vector, so that its particular bFunc has   *)
(*   fewer arguments, hence fewer derivatives and fewer undetermined        *)
(*   coefficients that must be solved.                                      *)
(* The routine then iterates, so that a variable is also counted as         *)
(*   clearly backward-looking if the only other variables in its equation   *)
(*   are lagged variables, date t shocks, and other variables dated t       *)
(*   that have already been shown to be clearly backward-looking.           *)

AIMBackLookingVars[eqns_]:= AIMBackLookingVars[eqns] =
Module[{eqnsnoleads,symbols,nsv,nsvargs,possiblensv,nextpos},
eqnsnoleads = Select[Flatten[{eqns}],Cases[#,_[t+i_/;i>0],Infinity] =={}&] ;
symbols = Map[Join[AIMLagVars[#],Union[Cases[#,_[t],Infinity]]]&, eqnsnoleads] ;
{nsv,nsvargs} = FixedPoint[
	Function[nsv,
	  symbols = symbols /.Thread[nsv[[1]]->nsv[[2]]] ;
	  possiblensv = Map[Cases[#,_Symbol[t]]&, symbols] ;
	  nextpos = Flatten[Position[Map[Length,possiblensv], 1]] ;
	 {Join[nsv[[1]], Flatten[possiblensv[[nextpos]]]],
	   Join[nsv[[2]], DeleteCases[symbols[[nextpos]],_Symbol[t],Infinity]]}
	], {{},{}}] ;
{nsv /.x_[t]:>x, Map[Union[Flatten[#]]&,nsvargs]}
] ;

AIMBLVPos[eqns_]:= AIMBLVPos[eqns] =
Flatten[Map[Position[AIMVarNames[eqns],#]&, AIMBackLookingVars[eqns][[1]]]] ;


(* Calculate the steady state *)

AIMSS[eqns_,___,opts___Rule]:= AIMSS[eqns,AIMZeroTol] =
Module[{ssguess,precision,sseqns,symbols,ss},
 ssguess = AIMSSGuess /.{opts} /.AIMSSGuess->Table[0,{AIMNVars[eqns]}] ;
 precision = AIMPrecision /.{opts} /.AIMPrecision->Precision[eqns]-2 ;
 sseqns = Thread[(eqns /.Equal->Subtract /.AIMSSSubs)==0] ;
 AIMModelDiagnostics[eqns] ;
 Print["Finding steady state, AIMSSGuess->",InputForm[ssguess]] ;
 symbols = Complement[Union[Cases[sseqns,_Symbol,Infinity]],
						Join[AIMSSVars[eqns],{E}]] ;
 If[Length[symbols]>0,Print["Warning: found symbols ",symbols," in equations"]];
 ss = Chop[FindRoot[sseqns, Transpose[{AIMSSVars[eqns],ssguess}],
		MaxIterations->5000, WorkingPrecision->precision], AIMZeroTol] ;
 If[Head[ss]===FindRoot, $Failed, ss]
] ;


(* Front end that does error-checking and formats the output *)

AIMSeries[eqns_,deg_Integer]:=
Module[{soln,const,argsubs},
 If[AIMSS[eqns,AIMZeroTol] === $Failed ||
	(soln=AIMSoln[eqns,deg,AIMZeroTol]) === $Failed, $Failed,
 Print["Formatting output, time is ", Date[]] ;
 const = AIMSSVars[eqns] /.AIMSS[eqns,AIMZeroTol] ;
 argsubs = Thread[AIMGenericArgs[eqns]->AIMStateVars[eqns]] ;
 Thread[Through[AIMVarNames[eqns][t]] == const + (soln /.argsubs)]
]] ;


(* Front end for Linear AIM *)

AIMSoln[eqns_,1,zerotol_]:= AIMSoln[eqns,1,zerotol] =
Module[{eqnseq0,allvars,hmat,epsmat,soln,cofb,s0inv,alllagvars},
 eqnseq0 = Flatten[{eqns}] /.Equal->Subtract ;
 allvars = Flatten[Map[Through[AIMVarNames[eqns] [t+#]]&,
			 Range[-Max[AIMMaxLag[eqns],1], AIMMaxLead[eqns]]]] ;
 hmat = Outer[D,eqnseq0,allvars] /.AIMSSSubs /.AIMSS[eqns,zerotol] ;
 epsmat = Outer[D,eqnseq0,AIMShocks[eqns]] /.AIMSSSubs /.AIMSS[eqns,zerotol] ;
 If[(soln=AIMLinearSoln[hmat,AIMMaxLead[eqns]]) === $Failed, $Failed,
 {cofb,s0inv} = soln ;
 alllagvars = allvars[[Range[Max[AIMMaxLag[eqns],1]*AIMNVars[eqns]]]] ;
 Chop[cofb .alllagvars - s0inv .epsmat .AIMShocks[eqns] /.
		Thread[AIMStateVars[eqns]->AIMGenericArgs[eqns]], zerotol]
]] ;


(* This is the heart of the program.  Derivatives are evaluated at steady    *)
(*   state, the expectation is taken, and coefficients are solved using      *)
(*   the method of undetermined coefficients.                                *)
(* As in the simple version, we solve a certainty-equivalent version of      *)
(*   the problem first.  In addition, for the clearly backward-looking       *)
(*   variables, we delete terms and coefficients that are clearly zero,      *)
(*   and thus solve only those coefficients that are in general nonzero.     *)

AIMSoln[eqns_,deg_Integer,zerotol_]:= AIMSoln[eqns,deg,zerotol] =
Module[{args,blvpos,bfuncargs,drvindxs,cedrvindxs,stdrvindxs,cecoeffs,
	stcoeffs,nextceTerms,nextstTerms,cesubs,cesystem,cesoln,dum,stsubs,
	stsystem,stsoln},
 args = AIMGenericArgs[eqns] ;
 blvpos = AIMBLVPos[eqns] ;
 bfuncargs = Table[args,{AIMNVars[eqns]}] ;
 bfuncargs[[blvpos]] = AIMBackLookingVars[eqns][[2]] /.
			Thread[AIMStateVars[eqns]->AIMGenericArgs[eqns]] ;
 drvindxs = Compositions[deg,AIMNArgs[eqns]] ;
  cedrvindxs = Select[drvindxs, #[[-1]]==0 &] ;
  stdrvindxs = Select[drvindxs, #[[-1]] >1 &] ;
 cecoeffs = Table[Unique[],{AIMNVars[eqns]},{Length[cedrvindxs]}] ;
  stcoeffs = Table[Unique[],{AIMNVars[eqns]},{Length[stdrvindxs]}] ;
 nextceTerms = cecoeffs .Map[Apply[Times,Power[args,#]]&, cedrvindxs] ;
  nextstTerms = stcoeffs .Map[Apply[Times,Power[args,#]]&, stdrvindxs] ;
 nextceTerms[[blvpos]] = Map[ nextceTerms[[#]] /. 
			Thread[Complement[args,bfuncargs[[#]]]->0]&, blvpos] ;
  nextstTerms[[blvpos]] = 0 ;
 cecoeffs[[blvpos]] = Map[Complement[Cases[nextceTerms[[#]],_Symbol,Infinity],
						bfuncargs[[#]]]&, blvpos] ;
  stcoeffs[[blvpos]] = {} ;
 cecoeffs = Flatten[cecoeffs] ;
  stcoeffs = Flatten[stcoeffs] ;
 cesubs = Thread[Map[bFunc,AIMVarNames[eqns]] ->
	MapThread[Apply[Function,{#1,#2}]&, {bfuncargs,
				AIMSoln[eqns,deg-1,zerotol] + nextceTerms}]] ;
 Print["Calculating CE Derivatives, time is ", Date[]] ;
 Print["Undetermined CE coefficients to solve: ",Length[cecoeffs]] ;
 cesoln = If[AIMNArgs[eqns]===1, PrependTo[args,dum]; nextceTerms,
   cesystem = Chop[Flatten[CoefficientArrays[AIMCEDerivatives[eqns,deg][0]
					/.cesubs, Drop[args,-2]]], zerotol] ;
   Print["Calculating CE solution, time is ", Date[]] ;
   nextceTerms /.Chop[Flatten[NSolve[cesystem, cecoeffs]], zerotol]] ;
 If[stcoeffs==={}, stsoln=0,
  stsubs = Thread[Map[bFunc,AIMVarNames[eqns]] ->
	MapThread[Apply[Function,{#1,#2}]&, {bfuncargs,
	AIMSoln[eqns,deg-1,zerotol] + cesoln + nextstTerms}]] ;
  Print["Calculating Stoch Derivatives, time is ", Date[]] ;
  stsystem = Chop[Expand[Flatten[CoefficientArrays[Chop[
	AIMStDerivatives[eqns,deg][0] /.stsubs, zerotol],Drop[args,-1]]]],
			zerotol] /.eps[x_][_]^n_->mom[x,n] /.eps[_][_]->0 ;
  Print["Undetermined Stoch coefficients to solve: ",Length[stcoeffs]] ;
  Print["Calculating Stoch solution, time is ", Date[]] ;
  stsoln = nextstTerms /.Chop[Flatten[NSolve[stsystem, stcoeffs]], zerotol] ;
 ] ;
 AIMSoln[eqns,deg-1,zerotol] + cesoln + stsoln
] ;

(* That's essentially it.  The following routine calculates derivatives	     *)
(*   of the equations composed with the (unknown) solution functions b.	     *)
(* Like the simple version, we use univariate differentiation to compute     *)
(*   the multivariate derivatives of the function Fob.  Unlike the simple    *)
(*   version, we compute the certainty-equivalent derivatives (those not     *)
(*   involving Sigma) first, and the stochastic derivatives (those involving *)
(*   Sigma) second.  This allows us to avoid ever computing the first        *)
(*   derivatives of Fob with respect to Sigma, which we know from theory     *)
(*   must be zero.                                                           *)

AIMCEDerivatives[eqns_,2]:= AIMCEDerivatives[eqns,2] =
Derivative[2][Function[tAIM, Evaluate[AIMSubBFuncsIntoEqns[eqns] /.
 Thread[AIMStateVars[eqns]->tAIM*Join[Drop[AIMGenericArgs[eqns],-2],{1,0}]]]]] ;

AIMCEDerivatives[eqns_,deg_Integer]:= AIMCEDerivatives[eqns,deg] =
AIMCEDerivatives[eqns,deg-1]' ;


AIMStDerivatives[eqns_,2]:= AIMStDerivatives[eqns,2] =
Function[tAIM, Evaluate[Derivative[0,2] [Function[{tAIM1,tAIM2},
	Evaluate[AIMSubBFuncsIntoEqns[eqns] /.Thread[AIMStateVars[eqns]->
	Append[tAIM1*Drop[AIMGenericArgs[eqns],-1],tAIM2]]]]] [tAIM,tAIM]]] ;

AIMStDerivatives[eqns_,deg_Integer]:= AIMStDerivatives[eqns,deg] =
AIMStDerivatives[eqns,deg-1]' ;


AIMSubBFuncsIntoEqns[eqns_]:= AIMSubBFuncsIntoEqns[eqns] =
With[{deveqns = eqns /.x_Symbol[t+i_.]:>Symbol[SymbolName[x]<>"AIMSS"] +x[t+i] 
				/.Equal->Subtract /.AIMSS[eqns,AIMZeroTol]},
 AIMSubOutLeadVars[eqns,deveqns,AIMMaxLead[eqns]]
] ;


AIMBFuncs[eqns_]:= AIMBFuncs[eqns] =
Module[{bfuncs},
 bfuncs = Map[Apply[bFunc[#],AIMStateVars[eqns]]&, AIMVarNames[eqns]] ;
 bfuncs[[AIMBLVPos[eqns]]] = MapThread[Apply[bFunc[#1],#2]&,
						AIMBackLookingVars[eqns]] ;
 bfuncs
] ;


AIMSubOutLeadVars[origeqns_,eqnssofar_,0]:= 
eqnssofar /.Thread[Through[AIMVarNames[origeqns][t]]->AIMBFuncs[origeqns]] /.
					eps[x_][t+i_]->Sigma*eps[x][t+i] ;

AIMSubOutLeadVars[origeqns_,eqnssofar_,lead_Integer]:=
With[{reducedeqns=eqnssofar /.Thread[Through[AIMVarNames[origeqns][t+lead]]->
					(AIMBFuncs[origeqns]/.t->t+lead)]},
 AIMSubOutLeadVars[origeqns, reducedeqns, lead-1]
] ;


(* Print out model diagnostics (obviously) *)

AIMModelDiagnostics[eqns_] := (
Print["\nModel Diagnostics:"] ;
Print["Number of equations: ",AIMNEqns[eqns]] ;
Print["Number of variables: ",AIMNVars[eqns]] ;
Print["Number of shocks:    ",Length[AIMShocks[eqns]]] ;
Print["Maximum lag:         ",AIMMaxLag[eqns]] ;
Print["Maximum lead:        ",AIMMaxLead[eqns]] ;
Print["Lagged variables: ",AIMLagVars[eqns]] ;
Print["Shocks: ",AIMShocks[eqns]] ;
Print[" together with Sigma, these yield ",AIMNArgs[eqns]," state variables\n"];
Print["Variables ",AIMBackLookingVars[eqns][[1]]," are clearly backward-looking"];
Print[" using the following reduced state variable sets for these variables:"] ;
Print[AIMBackLookingVars[eqns][[2]]] ;
Print["List of all variables: ",AIMVarNames[eqns]] ;
Print["Treating steady state and final coeff values < ",
				N[AIMZeroTol], " as zero (AIMZeroTol)\n"] ;
) ;


(* Everything that follows is Linear AIM.  See Anderson and Moore (1985)    *)
(*   for a description of the algorithm.                                    *)
(* In contrast to the simple version of the code, here we perform "exact"   *)
(*   ShiftRights first (i.e., shift equations forward that have a complete  *)
(*   row of zeros in the lead matrix, which reduces the number of calls to  *)
(*   the singular value decomposition and thus reduces potential numerical  *)
(*   problems.)                                                             *)
(* Also, we compute the stability conditions of the model using the Schur   *)
(*   decomposition rather than eigenvectors.  This should be more           *)
(*   numerically robust than eigenvectors, and the Schur decomposition is   *)
(*   guaranteed to exist for all models, while linearly independent         *)
(*   eigenvectors are not.  The disadvantage of using Schur vectors is that *)
(*   computing them is slow: first, the computation is inherently slower,   *)
(*   and second, the fact that the code must swap Schur blocks using a      *)
(*   routine that I (ets) wrote, rather than internal optimized machine     *)
(*   code, slows it down quite a bit.                                       *)

AIMLinearSoln[hmat_,nleads_Integer] :=
Module[{hrows,hcols,shiftedh,qmat,stabconds,bmat,smat},
 {hrows,hcols} = Dimensions[hmat] ;
 {shiftedh,qmat} = NestWhile[AIMExactShiftRight, {hmat,{}},
				Min[Map[Max,Abs[AIMLastBlock[#[[1]]]]]] ==0 &,
 				1, nleads] ;
 {shiftedh,qmat} = NestWhile[AIMShiftRight, {shiftedh,qmat},
 				MatrixRank[AIMLastBlock[#[[1]]]] <hrows &,
 				1, nleads*hrows] ;
 Print["Lead matrix is full rank; computing stability conditions"] ;
 stabconds = AIMLinearStabConds[Transpose[AIMAR1Form[shiftedh]]] ;
 Print["Model has ",Length[stabconds], " unstable roots"] ;
 qmat = Join[qmat,stabconds] ;
 If[Length[qmat]<hrows*nleads, Print["Multiple Linear Solutions"]; $Failed,
 If[Length[qmat]>hrows*nleads || MatrixRank[AIMLastBlock[qmat]]<hrows*nleads,
					Print["No Linear Solutions"]; $Failed,
 bmat = LinearSolve[AIMLastBlock[qmat],-AIMFirstBlocks[qmat]] ;
 smat = hmat .BlockMatrix[{{IdentityMatrix[hcols-hrows*nleads]},
				{ZeroMatrix[hrows*nleads,hrows],bmat}}] ;
 Chop[{Take[bmat,hrows], Inverse[AIMLastBlock[smat]]}, AIMZeroTol]
]]] ;


AIMLinearSoln[hmat_,0] :=
With[{cofb=LinearSolve[AIMLastBlock[hmat],-AIMFirstBlocks[hmat]]},
 Chop[{cofb, Inverse[AIMLastBlock[hmat]]}, AIMZeroTol]
] ;


AIMExactShiftRight[{hmatold_,qmatsofar_}]:=
Module[{hrows=Length[hmatold],zerorows,firstblocks},
 zerorows = Position[Map[Max, Abs[AIMLastBlock[hmatold]]], 0] ;
 firstblocks = AIMFirstBlocks[hmatold] ;
 Print["Shifting ",Length[zerorows]," equations ",zerorows," forward"];
 {ReplacePart[hmatold, BlockMatrix[{{ZeroMatrix[hrows,hrows],firstblocks}}],
		  					zerorows,zerorows],
   Join[qmatsofar,firstblocks[[Flatten[zerorows]]]]}
] ;


AIMShiftRight[{hmatold_,qmatsofar_}]:=
Module[{hrows=Length[hmatold],svdu,svdsig,hmatnew,zerorows,firstblocks},
 {svdu,svdsig} = Take[SingularValueDecomposition[AIMLastBlock[hmatold]],2] ;
 hmatnew = Transpose[svdu] .hmatold ;
 zerorows = Map[List,Range[MatrixRank[svdsig]+1,hrows]] ;
Print["Shifting ",Length[zerorows]," linear combinations of equations forward"];
 firstblocks = AIMFirstBlocks[hmatnew] ;
 Chop[{ReplacePart[hmatnew,BlockMatrix[{{ZeroMatrix[hrows,hrows],firstblocks}}],
 							 zerorows,zerorows],
   Join[qmatsofar,firstblocks[[Flatten[zerorows]]]]}, AIMZeroTol]
] ;


AIMAR1Form[shiftedh_]:= AIMAR1Form[shiftedh] =
With[{hrows=Length[shiftedh],hcols=Length[shiftedh[[1]]]},
 Chop[BlockMatrix[{{ZeroMatrix[hcols-2*hrows,hrows],
						IdentityMatrix[hcols-2*hrows]},
	{LinearSolve[AIMLastBlock[shiftedh],-AIMFirstBlocks[shiftedh]]}}],
	AIMZeroTol]
] ;


AIMLastBlock[matrix_]:=
With[{dims=Dimensions[matrix]},
 SubMatrix[matrix,{1,dims[[2]]-dims[[1]]+1},{dims[[1]],dims[[1]]}]
] ;


AIMFirstBlocks[matrix_]:=
With[{dims=Dimensions[matrix]},
 SubMatrix[matrix,{1,1},{dims[[1]],dims[[2]]-dims[[1]]}]
] ;


(* Mathematica's implementation of the Schur decomposition does not allow   *)
(*   the user to specify any particular order for the eigenvalues to        *)
(*   appear along the diagonal.  To compute the stability conditions of     *)
(*   the model, we require all eigenvalues greater than one in absolute     *)
(*   value to appear in the top part of the diagonal.  The following        *)
(*   two routines swap adjacent Schur blocks until this ordering is         *)
(*   achieved.                                                              *)
(* The routine is based on a 2002 paper by Jan Brandts entitled "Matlab     *)
(*   Code for Sorted Real Schur Forms," available at                        *)
(*   http://staff.science.uva.nl/~brandts.  One advantage of this           *)
(*   particular algorithm is that it allows for swapping 1x1 and 2x2        *)
(*   Schur blocks, so that we can avoid complex numbers entirely.  This     *)
(*   keeps annoying (albeit tiny) imaginary numbers from cropping up in     *)
(*   our stability conditions and hence our solution.                       *)
(* One can also use a Generalized Schur Decomposition instead of the        *)
(*   plain vanilla Schur Decomposition, but in Linear AIM the gains from    *)
(*   this are small because of the efficient way that AIM handles leads     *)
(*   (using ShiftRights to render the lead matrix invertible, and avoiding  *)
(*   stacking equations into the more numerically unstable AR1 form until   *)
(*   the very last minute).                                                 *)

AIMLinearStabConds[matrix_]:=
Module[{q,t,n,subdiagelts,twoblockpos,blockpos,eigvals,bigrtpos,swaplist},
 {q,t} = Chop[SchurDecomposition[matrix], AIMZeroTol] ;
 n = Length[matrix] ;
 subdiagelts = Map[t[[#+1,#]]&, Range[n-1]] ;
 twoblockpos = Flatten[Position[subdiagelts, i_/;i!=0]] ;
 blockpos = Complement[Range[n+1], twoblockpos+1] ;
 eigvals = Map[t[[#,#]]&, Drop[blockpos,-1]] ;
 eigvals[[Flatten[Map[Position[blockpos,#]&, twoblockpos]]]] +=
 			 Map[Sqrt[t[[#,#+1]] *t[[#+1,#]]]&, twoblockpos] ;
 bigrtpos = Flatten[Position[eigvals, i_/;Abs[i]>1+AIMZeroTol]] ;
 swaplist = Flatten[Map[Range[bigrtpos[[#]]-1,#,-1]&, Range[Length[bigrtpos]]]];
 {q,t,blockpos} = Fold[AIMSwapSchurBlocks, {Transpose[q],t,blockpos}, swaplist];
 Take[q, blockpos[[Length[bigrtpos]+1]] -1]
] ;


AIMSwapSchurBlocks[{q_,t_,blockpos_},swappos_]:=
Module[{index1,index2,n1,n2,a11,a12,a22,xmat,ans,xq,index12,qswpd,tprime,tswpd,
	blockposswpd},
 index1 = Range[blockpos[[swappos]], blockpos[[swappos+1]]-1] ;
 index2 = Range[blockpos[[swappos+1]], blockpos[[swappos+2]]-1] ;
 n1 = Length[index1] ;
 n2 = Length[index2] ;
 a11 = Take[t, index1, index1] ;
 a12 = Take[t, index1, index2] ;
 a22 = Take[t, index2, index2] ;
 xmat = Table[Unique[x],{n1},{n2}] ;
 ans = Flatten[NSolve[Flatten[a11.xmat-xmat.a22]-Flatten[a12], Flatten[xmat]]];
 xq = First[QRDecomposition[ BlockMatrix[{{-xmat /.ans, IdentityMatrix[n1]},
 				  {IdentityMatrix[n2], ZeroMatrix[n2,n1]}}] ]] ;
 index12 = Join[index1,index2] ;
 qswpd = q ;                                         (* premultiply q by xq *)
  qswpd[[index12]] = xq .qswpd[[index12]] ;
 tprime = Transpose[t] ;                       (* transform t to xq .t .xq' *)
  tprime[[index12]] = xq .tprime[[index12]] ;
  tswpd = Transpose[tprime] ;
  tswpd[[index12]] = xq .tswpd[[index12]] ;
 blockposswpd = ReplacePart[blockpos, blockpos[[swappos+1]] +n2-n1, swappos+1] ;
 Chop[{qswpd,tswpd,blockposswpd}, AIMZeroTol]
] ;