Source

OCaml Map Reduce Project / apps / nbody / nbody.ml

open Util

(* Create a transcript of body positions for `steps` time steps *)
let make_transcript (bodies : (string * body) list) (steps : int) : string = 
  let bodies_ref = ref bodies in
  let str = ref "" in
  let nbody_map_reduce () =
    let to_string_tuple (id, b) = (id, marshal b) in
    let kv_pairs = List.map to_string_tuple !bodies_ref in
    let mapped =
      Map_reduce.map kv_pairs (marshal !bodies_ref) "apps/nbody/mapper.ml" in
    let combined = Map_reduce.combine mapped in
    let reduced =
      Map_reduce.reduce combined "" "apps/nbody/reducer.ml" in
    let unwrap (k, v) = 
      match v with
      | [] -> failwith "no body was ever recovered"
      | h::_ -> (k, unmarshal h) in
    List.map unwrap reduced in
  str := string_of_bodies bodies;
  for i = 1 to steps do
    (bodies_ref := nbody_map_reduce (); 
    str := string_of_bodies !bodies_ref) done;
  !str

let simulation_of_string = function
  | "binary_star" -> Simulations.binary_star
  | "diamond" -> Simulations.diamond
  | "orbit" -> Simulations.orbit
  | "swarm" -> Simulations.swarm
  | "system" -> Simulations.system
  | "terrible_situation" -> Simulations.terrible_situation
  | "zardoz" -> Simulations.zardoz
  | _ -> failwith "Invalid simulation name. Check `shared/simulations.ml`"

let main (args : string array) : unit = 
  if Array.length args < 3 then 
    print_endline "Usage: nbody <simulation> [<outfile>]
  <simulation> is the name of a simulation from shared/simulations.ml
  Results will be written to [<outfile>] or stdout."
  else begin
    let (num_bodies_str, bodies) = simulation_of_string args.(2) in
    let transcript = make_transcript bodies 1 in
    let out_channel = 
      if Array.length args > 3 then open_out args.(3) else stdout in
    output_string out_channel (num_bodies_str ^ "\n" ^ transcript);
    close_out out_channel end