Working on new constraint algorithm

libraries/quern/README

+Quern 0.9 (beta)
+
+Public domain code for sparse up-looking Givens QR, including incomplete
+drop tolerance version for preconditioning iterative solvers.
+
+Edit Makefile to fit your set up (in particular, MATLAB-specific stuff).
+Currently it's set up for debugging; change the flags to get optimized
+versions of the code.
+
+In addition to the library libquern.a (with header quern.h), there are
+two mex files (quern_internal_mex and quern_cgnr_mex) as well as a
+stand-alone executable test_mat for running tests on matrices stored
+in .mat files.
+
+- Robert Bridson
+  UBC Computer Science
+  February 4, 2009

libraries/quern/include/quern.h

+/*
+* Quern: a sparse QR library.
+* This code is in the public domain.
+* - Robert Bridson
+*/
+
+#ifndef QUERN_H
+#define QUERN_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Function return values: nonzero indicates an error
+#define QUERN_OK 0
+#define QUERN_INPUT_ERROR 1
+#define QUERN_OUT_OF_MEMORY 2
+
+// Takes a compressed sparse column format matrix and outputs it in compressed
+// sparse row format (suitable for the QR factorization routines).
+// Requires A1_row_start to have room for num_columns+1 entries,
+//          A1_column_index to have room for all nonzeros,
+//      and A1_value to have room for all nonzeros.
+int QUERN_convert_column_format_to_row_format(int num_rows,
+                                              int num_columns,
+                                              const int* A0_column_start,
+                                              const int* A0_row_index,
+                                              const double* A0_value,
+                                              int* A1_row_start,
+                                              int* A1_column_index,
+                                              double* A1_value);
+
+// Takes a compressed sparse row format m*n matrix and finds a reversed
+// breadth-first search column ordering, suitable for incomplete QR
+// preconditioning.
+// Requires column_order to have room for n entries.
+int QUERN_get_rbfs_column_ordering(int m,
+                                   int n,
+                                   const int* A_row_start,
+                                   const int* A_column_index,
+                                   int* column_order);
+
+// Takes a length n column ordering and a compressed sparse row format m*n
+// matrix, then reorders each row accordingly.
+int QUERN_reorder_columns(int m,
+                          int n,
+                          const int* column_order,
+                          const int* A_row_start,
+                          int* A_column_index,
+                          double* A_value);
+
+// Takes a compressed sparse row format m*n matrix and finds a row ordering
+// which minimizes the lower profile---in particular, if A's rows can be
+// permuted to make it upper triangular, row_order will do the trick.
+// This is recommended as a cheap pre-process before QR, at least if nothing
+// is known about the current structure, to improve performance.
+// Requires row_order to have room for m integers.
+int QUERN_get_profile_row_ordering(int m,
+                                   int n,
+                                   const int* A_row_start,
+                                   const int* A_column_index,
+                                   int* row_order);
+
+// Reorders a given input vector x into an output vector y of length m.
+int QUERN_reorder_vector(int m,
+                         const int* order,
+                         const double* x,
+                         double* y);
+
+// Applies the inverse order to a given input vector x, saving in an output
+// vector y of length m.
+int QUERN_inverse_order_vector(int m,
+                               const int* order,
+                               const double* x,
+                               double* y);
+
+// Takes a compressed sparse row format m*n matrix (m>=n) and outputs the upper
+// triangular R factor from QR, also in compressed sparse row format. If
+// row_order is non-null, it should contain the order in which rows of A should be
+// taken; if it is null, the natural ordering is used. Memory for R is
+// allocated internally; use QUERN_free_result to free.
+int QUERN_compute_qr_without_q(int m,
+                               int n,
+                               const int* A_row_start,
+                               const int* A_column_index,
+                               const double* A_value,
+                               const int* row_order,
+                               int** ptr_R_row_start,
+                               int** ptr_R_column_index,
+                               double** ptr_R_value);
+
+// Takes a compressed sparse row format m*n matrix (m>=n) and outputs the QR
+// factors. The storage for Q encodes a sequence of Givens rotations and row
+// swaps; it is not directly usable as a standard sparse matrix. R, however,
+// is stored in standard compressed sparse row format. If row_order is non-null
+// it should contain the order in which rows of A should be taken; if it is
+// null, the natural ordering is used. Memory for Q and R is allocated
+// internally; use QUERN_free_result to free.
+int QUERN_compute_qr(int m,
+                     int n,
+                     const int* A_row_start,
+                     const int* A_column_index,
+                     const double* A_value,
+                     const int* row_order,
+                     int** ptr_Q_row_start,
+                     int** ptr_Q_column_index,
+                     double** ptr_Q_value,
+                     int** ptr_R_row_start,
+                     int** ptr_R_column_index,
+                     double** ptr_R_value);
+
+// Takes a compressed sparse row format m*n matrix (m>=n) and outputs the upper
+// triangular R factor from QR, with small entries dropped, also in compressed
+// sparse row format. If row_order is non-null, it should contain the order in
+// which rows of A should be taken; if it is null, the natural ordering is used.
+// Memory for R is allocated internally; use QUERN_free_result to free.
+int QUERN_compute_incomplete_qr_without_q(int m,
+                                          int n,
+                                          const int* A_row_start,
+                                          const int* A_column_index,
+                                          const double* A_value,
+                                          const int* row_order,
+                                          double drop_tolerance,
+                                          int** ptr_R_row_start,
+                                          int** ptr_R_column_index,
+                                          double** ptr_R_value);
+
+// Takes a compressed sparse row format m*n matrix (m>=n) and outputs incomplete
+// QR factors. The storage for Q encodes a sequence of Givens rotations and row
+// swaps; it is not directly usable as a standard sparse matrix. R, however,
+// is stored in standard compressed sparse row format. If row_order is non-null
+// it should contain the order in which rows of A should be taken; if it is
+// null, the natural ordering is used. Memory for Q and R is allocated
+// internally; use QUERN_free_result to free.
+int QUERN_compute_incomplete_qr(int m,
+                                int n,
+                                const int* A_row_start,
+                                const int* A_column_index,
+                                const double* A_value,
+                                const int* row_order,
+                                double drop_tolerance,
+                                int** ptr_Q_row_start,
+                                int** ptr_Q_column_index,
+                                double** ptr_Q_value,
+                                int** ptr_R_row_start,
+                                int** ptr_R_column_index,
+                                double** ptr_R_value);
+
+// Free the memory allocated during QR factorization (for either Q or R).
+// After calling this, do not try to access the factor again!
+void QUERN_free_result(int* row_start,
+                       int* column_index,
+                       double* value);
+
+// Compute the product of an m*n CSR matrix with n-vector input in
+// m-vector result.
+int QUERN_multiply(int m,
+                   int n,
+                   const int* row_start,
+                   const int* column_index,
+                   const double* value,
+                   const double* input,
+                   double* result);
+
+// Compute the product of an m*n CSR matrix transposes with m-vector input in
+// n-vector result.
+int QUERN_multiply_transpose(int m,
+                             int n,
+                             const int* row_start,
+                             const int* column_index,
+                             const double* value,
+                             const double* input,
+                             double* result);
+
+// Compute the product Q*x in-place (Q from QR factorization).
+// Note: if the input is naturally a length n vector, you should pad it with
+// zeros to make it length m.
+// If a row reordering of A was involved, you probably want to
+// call QUERN_inverse_order_vector afterwards.
+int QUERN_multiply_with_q(int m,
+                          const int* Q_row_start,
+                          const int* Q_column_index,
+                          const double* Q_value,
+                          double* x);
+
+// Compute the product Q^T*x in-place (Q from QR factorization).
+// If a row reordering of A was involved, you probably want to
+// call QUERN_reorder_vector first.
+int QUERN_multiply_with_q_transpose(int m,
+                                    const int* Q_row_start,
+                                    const int* Q_column_index,
+                                    const double* Q_value,
+                                    double* x);
+
+// Compute result=R^{-1}*rhs (R from QR factorization).
+// The vector result may actually be aliased to rhs, for an in-place solve.
+int QUERN_solve_with_r(int n,
+                       const int* R_row_start,
+                       const int* R_column_index,
+                       const double* R_value,
+                       const double* rhs,
+                       double* result);
+
+// Compute R^{-T}*x in-place (R from QR factorization).
+int QUERN_solve_with_r_transpose_in_place(int n,
+                                          const int* R_row_start,
+                                          const int* R_column_index,
+                                          const double* R_value,
+                                          double* x);
+
+// CGNR solver for A^T*A*x=rhs, starting with zero initial guess,
+// with R as preconditioner.
+int QUERN_solve_with_CGNR(int m,
+                          int n,
+                          const int* A_row_start,
+                          const int* A_column_index,
+                          const double* A_value,
+                          const double* rhs,
+                          const int* R_row_start,
+                          const int* R_column_index,
+                          const double* R_value,
+                          int max_iterations,
+                          double absolute_convergence_tolerance,
+                          double* x,
+                          int* return_solved,
+                          int* return_iterations,
+                          double* return_residual_norm);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

libraries/quern/include/quern_list.h

+#ifndef QUERN_LIST_H
+#define QUERN_LIST_H
+
+#include <algorithm>
+#include <cstdlib>
+
+// This is a light-weight singly linked list with half-decent
+// allocation, suitable only for plain-old-data, and designed
+// to make it easy to use robustly from plain C (e.g. no
+// exceptions, basic memory management with malloc/free and
+// returned error codes).
+
+namespace Quern {
+
+template<typename T>
+struct ListEntry
+{
+   T data;
+   ListEntry<T>* next;
+
+   ListEntry(void) {}
+   ListEntry(const T& data_, ListEntry<T>* next_=0)
+      : data(data_), next(next_) {}
+};
+
+const int puddle_size=1020;
+
+template<typename T>
+struct Puddle
+{
+   int end;
+   Puddle* next;
+   ListEntry<T> array[puddle_size];
+};
+
+template<typename T>
+struct Pool
+{
+   Puddle<T>* puddle_head;
+   ListEntry<T>* free_head;
+
+   Pool(void) : puddle_head(0), free_head(0) {}
+
+   ~Pool(void)
+   {
+      Puddle<T>* p=puddle_head;
+      while(p){
+         Puddle<T>* p_next=p->next;
+         std::free(p);
+         p=p_next;
+      }
+      puddle_head=0;
+      free_head=0;
+   }
+
+   ListEntry<T>* allocate(void)
+   {
+      // if we just have something free at hand, use it
+      if(free_head){
+         ListEntry<T>* p=free_head;
+         free_head=free_head->next;
+         return p;
+      }
+      // otherwise, if we know we need another puddle in the pool, try for it
+      if(!puddle_head || puddle_head->end==puddle_size){
+         Puddle<T>* new_puddle=(Puddle<T>*)std::malloc(sizeof(Puddle<T>));
+         if(!new_puddle) return 0;
+         new_puddle->end=0;
+         new_puddle->next=puddle_head;
+         puddle_head=new_puddle;
+      }
+      return &puddle_head->array[puddle_head->end++];
+   }
+
+   void release(ListEntry<T>* entry)
+   {
+      entry->next=free_head;
+      free_head=entry;
+   }
+};
+
+template<typename T>
+struct ListIterator
+{
+   ListEntry<T>** handle; // pointer to a "next" or "head" pointer in the list
+
+   ListIterator(ListEntry<T>** handle_=0)
+      : handle(handle_)
+   {}
+
+   bool still_going(void) const
+   {
+      assert(handle);
+      return *handle!=0;
+   }
+
+   void operator++(void)
+   {
+      assert(handle);
+      assert(*handle);
+      handle=&((*handle)->next);
+   }
+
+   T& operator*(void)
+   {
+      assert(handle);
+      assert(*handle);
+      return (*handle)->data;
+   }
+
+   T* operator->(void)
+   {
+      assert(handle);
+      assert(*handle);
+      return &(*handle)->data;
+   }
+
+   const T* operator->(void) const
+   {
+      assert(handle);
+      assert(*handle);
+      return &(*handle)->data;
+   }
+
+   const T& operator*(void) const
+   {
+      assert(handle);
+      assert(*handle);
+      return (*handle)->data;
+   }
+};
+
+template<typename T>
+struct List
+{
+   Pool<T>* pool;
+   ListEntry<T>* head;
+   int length;
+
+   List(void) {}
+
+   void init(Pool<T>* pool_)
+   {
+      pool=pool_;
+      head=0;
+      length=0;
+   }
+
+   ListIterator<T> begin(void)
+   {
+      return ListIterator<T>(&head);
+   }
+
+   bool empty(void)
+   {
+      return head==0;
+   }
+
+   // p ends up referring to the next element
+   void erase(ListIterator<T>& p)
+   {
+      assert(p.handle);
+      assert(*p.handle);
+      ListEntry<T>* q=*p.handle;
+      *p.handle=q->next;
+      pool->release(q);
+      --length;
+   }
+
+   T& front(void)
+   {
+      assert(head);
+      return head->data;
+   }
+
+   const T& front(void) const
+   {
+      assert(head);
+      return head->data;
+   }
+
+   // put the new element before the one p initially refers to;
+   // if successful, p ends up referring to the new element
+   bool insert(ListIterator<T>& p, const T& data)
+   {
+      ListEntry<T>* q=pool->allocate();
+      if(!q) return false;
+      q->next=*p.handle;
+      q->data=data;
+      *p.handle=q;
+      ++length;
+      return true;
+   }
+
+   void pop_front(void)
+   {
+      assert(head);
+      ListEntry<T>* p=head;
+      head=p->next;
+      pool->release(p);
+      --length;
+   }
+
+   bool push_front(const T& entry)
+   {
+      ListEntry<T>* p=pool->allocate();
+      if(!p) return false;
+      p->data=entry;
+      p->next=head;
+      head=p;
+      ++length;
+      return true;
+   }
+
+   int size(void)
+   {
+      return length;
+   }
+
+   void swap(List<T>& other)
+   {
+      assert(pool==other.pool);
+      std::swap(head, other.head);
+      std::swap(length, other.length);
+   }
+};
+
+} // end namespace
+
+#endif

libraries/quern/src/quern.cpp

+#include <cassert>
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+#include "quern.h"
+
+int QUERN_convert_column_format_to_row_format(int num_rows,
+                                              int num_columns,
+                                              const int* A0_column_start,
+                                              const int* A0_row_index,
+                                              const double* A0_value,
+                                              int* A1_row_start,
+                                              int* A1_column_index,
+                                              double* A1_value)
+{
+   // check input
+   if(num_rows<=0 || num_columns<=0) return QUERN_INPUT_ERROR;
+   if(!A0_column_start || !A0_row_index || !A0_value) return QUERN_INPUT_ERROR;
+   if(!A1_row_start || !A1_column_index || !A1_value) return QUERN_INPUT_ERROR;
+   // figure out number of entries in each row
+   std::memset(A1_row_start, 0, (num_rows+1)*sizeof(int));
+   for(int i=0; i<num_columns; ++i){
+      if(A0_column_start[i]>A0_column_start[i+1]) return QUERN_INPUT_ERROR;
+      for(int j=A0_column_start[i]; j<A0_column_start[i+1]; ++j){
+         if(A0_row_index[j]<0 || A0_row_index[j]>=num_rows)
+            return QUERN_INPUT_ERROR;
+         ++A1_row_start[A0_row_index[j]+1];
+      }
+   }
+   // cumulative sum to get row_start
+   for(int i=0; i<num_rows; ++i)
+      A1_row_start[i+1]+=A1_row_start[i];
+   // use a temporary copy of row_start for keeping track of where we add
+   int* row_pointer=(int*)std::malloc(num_rows*sizeof(int));
+   if(!row_pointer) return QUERN_OUT_OF_MEMORY;
+   std::memcpy(row_pointer, A1_row_start, num_rows*sizeof(int));
+   // then fill in the entries
+   for(int i=0; i<num_columns; ++i){
+      for(int j=A0_column_start[i]; j<A0_column_start[i+1]; ++j){
+         int r=A0_row_index[j];
+         A1_column_index[row_pointer[r]]=i;
+         A1_value[row_pointer[r]]=A0_value[j];
+         ++row_pointer[r];
+      }
+   }
+   std::free(row_pointer);
+   return QUERN_OK;
+}
+
+int QUERN_multiply(int m,
+                   int n,
+                   const int* row_start,
+                   const int* column_index,
+                   const double* value,
+                   const double* input,
+                   double* result)
+{
+   if(m<=0 || n<=0 || !row_start || !column_index || !value
+         || !input || !result)
+      return QUERN_INPUT_ERROR;
+   int i, j;
+   double x;
+   for(i=0; i<m; ++i){
+      x=0;
+      for(j=row_start[i]; j<row_start[i+1]; ++j)
+         x+=value[j]*input[column_index[j]];
+      result[i]=x;
+   }
+   return QUERN_OK;
+}
+
+int QUERN_multiply_transpose(int m,
+                             int n,
+                             const int* row_start,
+                             const int* column_index,
+                             const double* value,
+                             const double* input,
+                             double* result)
+{
+   if(m<=0 || n<=0 || !row_start || !column_index || !value
+         || !input || !result)
+      return QUERN_INPUT_ERROR;
+   int i, j;
+   double x;
+   std::memset(result, 0, n*sizeof(double));
+   for(i=0; i<m; ++i){
+      x=input[i];
+      for(j=row_start[i]; j<row_start[i+1]; ++j)
+         result[column_index[j]]+=value[j]*x;
+   }
+   return QUERN_OK;
+}
+

libraries/quern/src/quern_order.cpp

+#include <cassert>
+#include <cstdlib>
+#include <cstring>
+#include <algorithm>
+#include "quern.h"
+
+int QUERN_get_rbfs_column_ordering(int m,
+                                   int n,
+                                   const int* A_row_start,
+                                   const int* A_column_index,
+                                   int* column_order)
+{
+   if(m<=0 || n<=0 || !A_row_start || !A_column_index || !column_order)
+      return QUERN_INPUT_ERROR;
+   // get some memory to work in
+   int* work=(int*)std::malloc((n+1+A_row_start[m]+n+n)*sizeof(int));
+   if(!work) return QUERN_OUT_OF_MEMORY;
+   // figure out number of entries in each column
+   int* column_start=work;
+   std::memset(column_start, 0, (n+1)*sizeof(int));
+   for(int i=0; i<m; ++i){
+      for(int j=A_row_start[i]; j<A_row_start[i+1]; ++j)
+         ++column_start[A_column_index[j]+1];
+   }
+   // cumulative sum to get column_start
+   for(int i=0; i<n; ++i)
+      column_start[i+1]+=column_start[i];
+   assert(column_start[n]==A_row_start[m]);
+   // list the columns now
+   int* row_index=column_start+(n+1);
+   int* column_pointer=row_index+column_start[n];
+   std::memcpy(column_pointer, column_start, n*sizeof(int));
+   for(int i=0; i<m; ++i){
+      for(int j=A_row_start[i]; j<A_row_start[i+1]; ++j){
+         int c=A_column_index[j];
+         row_index[column_pointer[c]++]=i;
+      }
+   }
+   // set up marker for BFS
+   char* column_marker=(char*)(column_pointer+n);
+   std::memset(column_marker, 0, n);
+   // and do as many BFS as we need to hit all connected components
+   int p=n;
+   for(int root=0; root<n; ++root) if(!column_marker[root]){
+      column_order[--p]=root;
+      column_marker[root]=1;
+      for(int i=p; i>=p; --i){
+         int j=column_order[i];
+         // add unmarked neighbour columns of j to ordering
+         for(int k=column_start[j]; k<column_start[j+1]; ++k){
+            int r=row_index[k];
+            for(int a=A_row_start[r]; a<A_row_start[r+1]; ++a){
+               int nbr=A_column_index[a];
+               if(!column_marker[nbr]){
+                  column_order[--p]=nbr;
+                  column_marker[nbr]=1;
+               }
+            }
+         }
+      }
+   }
+   assert(p==0);
+   std::free(work);
+   return QUERN_OK;
+}
+
+int QUERN_reorder_columns(int m,
+                          int n,
+                          const int* column_order,
+                          const int* A_row_start,
+                          int* A_column_index,
+                          double* A_value)
+{
+   if(m<=0 || n<=0 || !column_order
+         || !A_row_start || !A_column_index || !A_value)
+      return QUERN_INPUT_ERROR;
+
+   // Since we don't expect to have lots of empty columns, it would
+   // probably be superior to do this as two transposes --- but we'll
+   // worry about that optimization later.
+
+   // get some temporary storage for permuting and sorting
+   std::pair<int,double>* row=(std::pair<int,double>*)
+                                  std::malloc(n*sizeof(std::pair<int,double>));
+   if(!row) return QUERN_OUT_OF_MEMORY;
+   // and invert the permutation
+   int* inv=(int*)std::malloc(n*sizeof(int));
+   if(!inv){
+      std::free(row);
+      return QUERN_OUT_OF_MEMORY;
+   }
+   for(int i=0; i<n; ++i) inv[column_order[i]]=i;
+   // do it row by row
+   int k;
+   for(int i=0; i<m; ++i){
+      k=0;
+      for(int j=A_row_start[i]; j<A_row_start[i+1]; ++j)
+         row[k++]=std::make_pair(inv[A_column_index[j]], A_value[j]);
+      std::sort(row, row+k); // at some point should replace with radix-sort
+      k=0;
+      for(int j=A_row_start[i]; j<A_row_start[i+1]; ++j){
+         A_column_index[j]=row[k].first;
+         A_value[j]=row[k].second;
+         ++k;
+      } 
+   }
+   std::free(row);
+   std::free(inv);
+   return QUERN_OK;
+}
+
+int QUERN_get_profile_row_ordering(int m,
+                                   int n,
+                                   const int* A_row_start,
+                                   const int* A_column_index,
+                                   int* row_order)
+{
+   if(m<=0 || n<=0 || !A_row_start || !A_column_index || !row_order)
+      return QUERN_INPUT_ERROR;
+   // first count how many rows start at column i for each i=0, .., n-1
+   int* count=(int*)std::calloc(n+1, sizeof(int));
+   if(!count) return QUERN_OUT_OF_MEMORY;
+   for(int i=0; i<m; ++i)
+      if(A_row_start[i]<A_row_start[i+1])
+         ++count[A_column_index[A_row_start[i]]+1];
+   // then do a cumulative sum to find target locations of the rows
+   for(int j=2; j<n+1; ++j)
+      count[j]+=count[j-1];
+   // and now write out the row ordering with a second pass
+   for(int i=0; i<m; ++i){
+      if(A_row_start[i]<A_row_start[i+1])
+         row_order[count[A_column_index[A_row_start[i]]]++]=i;
+      else
+         row_order[count[n]++]=i;
+   }
+   std::free(count);
+   return QUERN_OK;
+}
+
+// Reorders a given input vector x into an output vector y of length m.
+int QUERN_reorder_vector(int m,
+                         const int* order,
+                         const double* x,
+                         double* y)
+{
+   if(m<=0 || !order || !x || !y) return QUERN_INPUT_ERROR;
+   for(int i=0; i<m; ++i)
+      y[i]=x[order[i]];
+   return QUERN_OK;
+}
+
+// Applies the inverse order to a given input vector x, saving in an output
+// vector y of length m.
+int QUERN_inverse_order_vector(int m,
+                               const int* order,
+                               const double* x,
+                               double* y)
+{
+   if(m<=0 || !order || !x || !y) return QUERN_INPUT_ERROR;
+   for(int i=0; i<m; ++i)
+      y[order[i]]=x[i];
+   return QUERN_OK;
+}

libraries/quern/src/quern_solver.cpp

+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+#include "quern.h"
+
+// TO-DO: use BLAS versions of these smaller kernels, if available
+
+double two_norm(int n, const double* x)
+{
+   double r=0;
+   for(int i=0; i<n; ++i) r+=x[i]*x[i];
+   return std::sqrt(r);
+}
+
+double two_norm_squared(int n, const double* x)
+{
+   double r=0;
+   for(int i=0; i<n; ++i) r+=x[i]*x[i];
+   return r;
+}
+
+// x=x+alpha*y
+void add_scaled(int n, double* x, double alpha, const double* y)
+{
+   for(int i=0; i<n; ++i) x[i]+=alpha*y[i];
+}
+
+// x=beta*x+y
+void scale_and_add(int n, double beta, double* x, const double* y)
+{
+   for(int i=0; i<n; ++i) x[i]=beta*x[i]+y[i];
+}
+
+int QUERN_solve_with_CGNR(int m,
+                          int n,
+                          const int* A_row_start,
+                          const int* A_column_index,
+                          const double* A_value,
+                          const double* rhs,
+                          const int* R_row_start,
+                          const int* R_column_index,
+                          const double* R_value,
+                          int max_iterations,
+                          double absolute_convergence_tolerance,
+                          double* x,
+                          int* return_solved,
+                          int* return_iterations,
+                          double* return_residual_norm)
+{
+   if(m<=0 || n<=0 || !A_row_start || !A_column_index || !A_value
+      || !rhs || !R_row_start || !R_column_index || !R_value || !x
+      || !return_solved || !return_iterations || !return_residual_norm)
+      return QUERN_INPUT_ERROR;
+   // default values
+   *return_solved=0;
+   *return_iterations=0;
+   *return_residual_norm=two_norm(n, rhs);
+   if(*return_residual_norm<=absolute_convergence_tolerance){
+      *return_solved=1;
+      return QUERN_OK;
+   }
+   // allocate some room to work in
+   double* working_vectors=(double*)std::malloc((3*n+m)*sizeof(double));
+   if(!working_vectors)
+      return QUERN_OUT_OF_MEMORY;
+   double* r=working_vectors;
+   double* s=r+n;
+   double* z=s+n;
+   double* u=z+n;
+   // set up CGNR
+   int check; 
+   std::memset(x, 0, n*sizeof(double));
+   std::memcpy(r, rhs, n*sizeof(double));
+   std::memcpy(u, rhs, n*sizeof(double));
+   check=QUERN_solve_with_r_transpose_in_place(n, R_row_start, R_column_index,
+                                               R_value, u);
+   if(check){ std::free(working_vectors); return check; }
+   check=QUERN_solve_with_r(n, R_row_start, R_column_index, R_value, u, z);
+   if(check){ std::free(working_vectors); return check; }
+   std::memcpy(s, z, n*sizeof(double));
+   double rho=two_norm_squared(n, u);
+   // the main loop
+   for(;;){
+      if(rho==0){ std::free(working_vectors); return QUERN_INPUT_ERROR; }
+      check=QUERN_multiply(m, n, A_row_start, A_column_index, A_value, s, u);
+      if(check){ std::free(working_vectors); return check; }
+      check=QUERN_multiply_transpose(m, n, A_row_start, A_column_index, A_value,
+                                     u, z);
+      if(check){ std::free(working_vectors); return check; }
+      double denom=two_norm_squared(m, u);
+      if(denom==0){ std::free(working_vectors); return QUERN_INPUT_ERROR; }
+      double alpha=rho/denom;
+      add_scaled(n, x, alpha, s);
+      add_scaled(n, r, -alpha, z);
+      ++*return_iterations;
+      *return_residual_norm=two_norm(n, r);
+      if(*return_residual_norm<=absolute_convergence_tolerance){
+         *return_solved=1;
+         break;
+      }
+      if(*return_iterations>max_iterations)
+         break;
+      std::memcpy(u, r, n*sizeof(double));
+      check=QUERN_solve_with_r_transpose_in_place(n, R_row_start,
+                                                  R_column_index, R_value, u);
+      if(check){ std::free(working_vectors); return check; }
+      check=QUERN_solve_with_r(n, R_row_start, R_column_index, R_value, u, z);
+      if(check){ std::free(working_vectors); return check; }
+      double rho_new=two_norm_squared(n, u);
+      double beta=rho_new/rho;
+      scale_and_add(n, beta, s, z);
+      rho=rho_new;
+   }
+   std::free(working_vectors);
+   return QUERN_OK;
+}

platforms/reference/src/ReferenceKernels.cpp

@@ -58,21 +58,21 @@
 using namespace OpenMM;
 using namespace std;
 
-int** allocateIntArray(int length, int width) {
+static int** allocateIntArray(int length, int width) {
     int** array = new int*[length];
     for (int i = 0; i < length; ++i)
         array[i] = new int[width];
     return array;
 }
 
-RealOpenMM** allocateRealArray(int length, int width) {
+static RealOpenMM** allocateRealArray(int length, int width) {
     RealOpenMM** array = new RealOpenMM*[length];
     for (int i = 0; i < length; ++i)
         array[i] = new RealOpenMM[width];
     return array;
 }
 
-int** copyToArray(const vector<vector<int> > vec) {
+static int** copyToArray(const vector<vector<int> > vec) {
     if (vec.size() == 0)
         return new int*[0];
     int** array = allocateIntArray(vec.size(), vec[0].size());
@@ -82,7 +82,7 @@ int** copyToArray(const vector<vector<int> > vec) {
     return array;
 }
 
-RealOpenMM** copyToArray(const vector<vector<double> > vec) {
+static RealOpenMM** copyToArray(const vector<vector<double> > vec) {
     if (vec.size() == 0)
         return new RealOpenMM*[0];
     RealOpenMM** array = allocateRealArray(vec.size(), vec[0].size());
@@ -92,7 +92,7 @@ RealOpenMM** copyToArray(const vector<vector<double> > vec) {
     return array;
 }
 
-void disposeIntArray(int** array, int size) {
+static void disposeIntArray(int** array, int size) {
     if (array) {
         for (int i = 0; i < size; ++i)
             delete[] array[i];
@@ -100,7 +100,7 @@ void disposeIntArray(int** array, int size) {
     }
 }
 
-void disposeRealArray(RealOpenMM** array, int size) {
+static void disposeRealArray(RealOpenMM** array, int size) {
     if (array) {
         for (int i = 0; i < size; ++i)
             delete[] array[i];
@@ -108,6 +108,20 @@ void disposeRealArray(RealOpenMM** array, int size) {
     }
 }
 
+static void findAnglesForShake(const System& system, vector<ReferenceRigidShakeAlgorithm::AngleInfo>& angles) {
+    for (int i = 0; i < system.getNumForces(); i++) {
+        const HarmonicAngleForce* force = dynamic_cast<const HarmonicAngleForce*>(&system.getForce(i));
+        if (force != NULL) {
+            for (int j = 0; j < force->getNumAngles(); j++) {
+                int atom1, atom2, atom3;
+                double angle, k;
+                force->getAngleParameters(j, atom1, atom2, atom3, angle, k);
+                angles.push_back(ReferenceRigidShakeAlgorithm::AngleInfo(atom1, atom2, atom3, angle));
+            }
+        }
+    }
+}
+
 void ReferenceInitializeForcesKernel::initialize(const System& system) {
 }
 
@@ -609,7 +623,9 @@ void ReferenceIntegrateVerletStepKernel::execute(OpenMMContextImpl& context, con
             delete constraints;
         }
         dynamics = new ReferenceVerletDynamics(context.getSystem().getNumParticles(), static_cast<RealOpenMM>(stepSize) );
-        constraints = new ReferenceRigidShakeAlgorithm(context.getSystem().getNumParticles(), numConstraints, constraintIndices, constraintDistances, masses, (RealOpenMM)integrator.getConstraintTolerance());
+        vector<ReferenceRigidShakeAlgorithm::AngleInfo> angles;
+        findAnglesForShake(context.getSystem(), angles);
+        constraints = new ReferenceRigidShakeAlgorithm(context.getSystem().getNumParticles(), numConstraints, constraintIndices, constraintDistances, masses, angles, (RealOpenMM)integrator.getConstraintTolerance());
         dynamics->setReferenceConstraintAlgorithm(constraints);
         prevStepSize = stepSize;
     }
@@ -668,7 +684,9 @@ void ReferenceIntegrateLangevinStepKernel::execute(OpenMMContextImpl& context, c
 				static_cast<RealOpenMM>(stepSize), 
 				static_cast<RealOpenMM>(tau), 
 				static_cast<RealOpenMM>(temperature) );
-        constraints = new ReferenceRigidShakeAlgorithm(context.getSystem().getNumParticles(), numConstraints, constraintIndices, constraintDistances, masses, (RealOpenMM)integrator.getConstraintTolerance());
+        vector<ReferenceRigidShakeAlgorithm::AngleInfo> angles;
+        findAnglesForShake(context.getSystem(), angles);
+        constraints = new ReferenceRigidShakeAlgorithm(context.getSystem().getNumParticles(), numConstraints, constraintIndices, constraintDistances, masses, angles, (RealOpenMM)integrator.getConstraintTolerance());
         dynamics->setReferenceConstraintAlgorithm(constraints);
         prevTemp = temperature;
         prevFriction = friction;
@@ -728,7 +746,9 @@ void ReferenceIntegrateBrownianStepKernel::execute(OpenMMContextImpl& context, c
 				static_cast<RealOpenMM>(stepSize), 
 				static_cast<RealOpenMM>(friction), 
 				static_cast<RealOpenMM>(temperature) );
-        constraints = new ReferenceRigidShakeAlgorithm(context.getSystem().getNumParticles(), numConstraints, constraintIndices, constraintDistances, masses, (RealOpenMM)integrator.getConstraintTolerance());
+        vector<ReferenceRigidShakeAlgorithm::AngleInfo> angles;
+        findAnglesForShake(context.getSystem(), angles);
+        constraints = new ReferenceRigidShakeAlgorithm(context.getSystem().getNumParticles(), numConstraints, constraintIndices, constraintDistances, masses, angles, (RealOpenMM)integrator.getConstraintTolerance());
         dynamics->setReferenceConstraintAlgorithm(constraints);
         prevTemp = temperature;
         prevFriction = friction;

platforms/reference/src/SimTKReference/ReferenceRigidShakeAlgorithm.h

@@ -26,6 +26,7 @@
 #define __ReferenceQShakeAlgorithm_H__
 
 #include "ReferenceConstraintAlgorithm.h"
+#include <utility>
 #include <vector>
 #include <set>
 
@@ -47,10 +48,12 @@ class ReferenceRigidShakeAlgorithm : public ReferenceConstraintAlgorithm {
       RealOpenMM* _distanceTolerance;
       RealOpenMM* _reducedMasses;
       bool _hasInitializedMasses;
-      std::vector<std::vector<int> > _rigidClusters;
-      std::vector<RealOpenMM**> _matrices;
+//      std::vector<std::vector<int> > _rigidClusters;
+//      std::vector<RealOpenMM**> _matrices;
+      std::vector<std::vector<std::pair<int, RealOpenMM> > > _matrix;
 
    public:
+      class AngleInfo;
 
       /**---------------------------------------------------------------------------------------
 
@@ -64,7 +67,7 @@ class ReferenceRigidShakeAlgorithm : public ReferenceConstraintAlgorithm {
 
          --------------------------------------------------------------------------------------- */
 
-      ReferenceRigidShakeAlgorithm( int numberOfAtoms, int numberOfConstraints, int** atomIndices, RealOpenMM* distance, RealOpenMM* masses, RealOpenMM tolerance );
+      ReferenceRigidShakeAlgorithm( int numberOfAtoms, int numberOfConstraints, int** atomIndices, RealOpenMM* distance, RealOpenMM* masses, std::vector<AngleInfo>& angles, RealOpenMM tolerance );
 
       /**---------------------------------------------------------------------------------------
 
@@ -171,6 +174,17 @@ class ReferenceRigidShakeAlgorithm : public ReferenceConstraintAlgorithm {
       int reportShake( int numberOfAtoms, RealOpenMM** atomCoordinates, std::stringstream& message );
 };
 
+class ReferenceRigidShakeAlgorithm::AngleInfo
+{
+public:
+    int atom1, atom2, atom3;
+    RealOpenMM angle;
+    AngleInfo(int atom1, int atom2, int atom3, RealOpenMM angle) :
+        atom1(atom1), atom2(atom2), atom3(atom3), angle(angle)
+    {
+    }
+};
+
 // ---------------------------------------------------------------------------------------
 
 #endif // __ReferenceQShakeAlgorithm_H__