The wave penetration project: energy dissipation by bottom friction

const int maxNonLinearLoopIter = 42;
std::vector<std::complex<double> > eta(mesh->_vertices.size(), 0.0);
// Picard iteration, non-linear loop
for (int iter = 0; iter < maxNonLinearLoopIter; ++iter)
{
	MildSlopeEquation equation;
	Solver<std::complex<double> > solver;
	solver.solve(equation.discretize(mesh, wave, R, cf, eta), mesh, eta);
}

(* 
	Computes the CRS "row ptr", "col ind" as well as a mapping from elementary matrices coefficient to "values" coefficient 
	ts: triangle (v1,v2,v3) list, nv: number of vertices
*)
let build_crs_profile ts nv =
	
	let rct = build_reverse_connectivity_table ts nv in

	(* Assembles a single row of the linear system *)
	let do_row ~init ~f row =
		rct.(row) |> List.fold ~init ~f:(fun acc ti ->
			let col0, col1, col2 = Int32.(to_int_exn ts.{ti,0}, to_int_exn ts.{ti,1}, to_int_exn ts.{ti,2}) in
			let trow = if col0 = row then 0 else if col1 = row then 1 else 2 in
			let do_col = f ti trow row 
			in acc |> do_col 0 col0 |> do_col 1 col1 |> do_col 2 col2
		)	
	in 

	(* Simulates the assembly in order to count the number of non zero coefficients *)
	let nnz = 
		(* color helps us to track already counted coeffs *)
		let color = Array.create ~len:(Array.length rct) (-1) in
		(* "fold" all the rows *)
		rct |> Array.foldi ~init:0 ~f:(fun row cnt _ ->
			row |> do_row ~init:cnt ~f:(fun ti trow row tcol col cnt ->
				(* Do not count coeffs twice *)
				if color.(col) < row then (
					color.(col) <- row;
					succ cnt
				) else
					cnt
			)
		)

	(* Perform the real assembly this time *)
	and nv = Array.length rct in		

	(* color helps us to track already counted coeffs *)
	let color = Array.create ~len:nv (-1)

	(* edges helps us to remember triangles linked to an edge *)
	and edges = Array.create ~len:nnz [|(0,0); (0,0)|]

	(* The result arrays, colidx will later be converted to a BigArray as BigArray does not provide an in-place sort function *)
	and rowptr = Array1.create int32 c_layout (nv + 1)
	and colidx = Array.create ~len:nnz Int32.zero
	and tt = Array2.create int32 c_layout (Array2.dim1 ts) 9 in

	(* This sorts the nz column idx of a row before updating rowptr, colidx and tt *)
	let commit_row rstart rend = 		
		Array.sort ~pos:rstart ~len:(rend - rstart) ~cmp:compare colidx;
		for pos = rstart to (rend-1) do
			let col = Int32.to_int_exn colidx.(pos) 
			and pos' = Int32.of_int_exn pos in
			let t0, idx0 = edges.(col).(0)
			and t1, idx1 = edges.(col).(1) in
			tt.{t0, idx0} <- pos';
			tt.{t1, idx1} <- pos'
		done
	in

	(* "fold" all the rows once again *)
 	let last = rct |> Array.foldi ~init:0 ~f:(fun row pos _ ->		
		let pos' = row |> do_row ~init:pos ~f:(fun ti trow row tcol col pos ->
			let idx = 3 * trow + tcol in
			if color.(col) < row then (
				(* We found the first edge going from row to col, 
					this could be the only one if it's a boundary edge *)
				color.(col) <- row;
				edges.(col).(0) <- (ti, idx);
				edges.(col).(1) <- (ti, idx);
				colidx.(pos) <- Int32.of_int_exn col;
				succ pos
			) else (
				(* there is a second edge going from row to col *)
				edges.(col).(1) <- (ti, idx);
				pos
			)
		) 
		in

		(* columns idx of the current row starts at position pos and ends at position pos' - 1 *)
		rowptr.{row} <- Int32.of_int_exn pos;
		commit_row pos pos';
		pos'
	)
	in rowptr.{nv} <- Int32.of_int_exn last;

	(rowptr, (Array1.of_array int32 c_layout colidx), tt)
;;

(* ts: triangle (v1,v2,v3) list, nv: number of vertices *)
let build_reverse_connectivity_table ts nv = 
	let rct = Array.create ~len:nv [] in
	List.iteri ts ~f:(fun t (v1,v2,v3) ->	
		rct.(v1) <- t::rct.(v1); 
		rct.(v2) <- t::rct.(v2); 
		rct.(v3) <- t::rct.(v3));
	rct
;;

let build_triangulation_dual_graph ts nv = 	
	let graph = Array.create ~len:(List.length ts) [] in
	let rct = build_reverse_connectivity_table ts nv in
	List.iteri ts ~f:(fun t (v1,v2,v3) ->
		graph.(t) <- List.dedup (rct.(v1) @ rct.(v2) @ rct.(v3)));
	graph
;;

let build_perfect_elemination_ordering graph =	
	let compare (p1,_) (p2,_) = compare p2 p1 in
	let heap = Heap.Removable.create ~min_size:(Array.length graph) ~cmp:compare () in
	let elems = Array.mapi (fun v _ -> Heap.Removable.add_removable heap (0,v)) graph in	
	let rec aux peo = 
		match Heap.Removable.pop heap with
		| Some (p,v) -> List.iter graph.(v) ~f:(fun v' -> 
				try					
					let (p',_) = Heap.Removable.Elt.value_exn elems.(v') in
					elems.(v') <- Heap.Removable.update heap elems.(v') (succ p', v')
				with
				| _ -> ();
			);
			aux (v::peo)
		| None -> List.rev peo
	in 
	aux []
;;

let do_greedy_coloring graph ordering =
	let n = Array.length graph in 
	let colors = Array.create ~len:n n in
	let perfect_color v =
		let rec aux color = function
			| c::cs' -> if color < c then color else aux (succ color) cs'
			| [] -> color
		in graph.(v) |> List.map ~f:(fun v -> colors.(v)) |> List.sort ~cmp:compare |> aux 0
	in
	List.iter ordering ~f:(fun v -> 	
		colors.(v) <- perfect_color v
	);
	colors
;;

let graph = build_triangulation_dual_graph ts (List.length vs)
let ordering = build_perfect_elemination_ordering graph
let perfect_coloring = do_greedy_coloring graph ordering

// Coloring computed by the greedy coloring algorithm
device_vector<int> colorsOnGpu(colors, &colors[nt]);

// tt is the mapping from local triangle element coefficient index
// to the index in the global linear system compressed row storage values
device_vector<int> ttOnGpu(tt, &tt[9 * nt]);

// Compressed row storage values on the GPU, initialized to 0 before summing
device_vector<float> result(rowptr[nv]);
fill(result.begin(), result.end(), 0.);

// How many colors ?
device_vector<int>::iterator maxIter = max_element(colorsOnGpu.begin(), colorsOnGpu.end());
int nbColor = 1 + *maxIter;
for (int color = 0; color < nbColor; ++color)
{
	// Only assemble triangles with that color
	Assemble<<<nbBlock, trianglesPerBlock>>>( 
	  coeffsOnGpu,
	  colorsOnGpu,
	  ttOnGpu,
	  result,
	  color
	);
}
copy(coeffsOnGpu.begin(), coeffsOnGpu.end(), coeffs);

__global__
  void Assemble(
    const KernelArray<float> coeffs, 
    const KernelArray<color_t> colors, 
    const KernelArray<int> tt, 
    KernelArray<float> result,
    color_t color) {
    const int myTriangleId = blockIdx.x * blockDim.x + threadIdx.x;
    if (myTriangleId < colors.size() && colors[myTriangleId] == color) {
      const int *mytt = &tt[9 * myTriangleId];
      const float *myCoeffs = &coeffs[9 * myTriangleId];
      result[mytt[0]] += myCoeffs[0];
      result[mytt[1]] += myCoeffs[1];
      result[mytt[2]] += myCoeffs[2];
      result[mytt[3]] += myCoeffs[3];
      result[mytt[4]] += myCoeffs[4];
      result[mytt[5]] += myCoeffs[5];
      result[mytt[6]] += myCoeffs[6];
      result[mytt[7]] += myCoeffs[7];
      result[mytt[8]] += myCoeffs[8];
    }
  }