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 rec nbody_map_reduce bodies steps =
    if steps = 0 then string_of_bodies bodies
    else
      let to_string_tuple (id, b) = (id, marshal b) in
      let kv_pairs = List.map to_string_tuple bodies in
      let mapped =
        Map_reduce.map kv_pairs (marshal bodies) "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::_ -> unmarshal h in
      string_of_bodies bodies ^
        (nbody_map_reduce (List.map unwrap reduced) (steps - 1)) in
  nbody_map_reduce bodies steps

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 60 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