Source

planets / fast_physics.ml

(*  Planets:  A celestial simulator
    Copyright (C) 2001-2003  Yaron M. Minsky

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*)

open StdLabels
open MoreLabels

open Printf
open State
open Constants
open Common

(* Faster physics implementation that is array-based *)
(* Only implements the act_all_on_all function. *)
(* Another useful thing to implement faster would be collision detection *)

(* all-float record, to let the compiler unbox in a loop *)
type fbody = { mutable x_pos: float;
               mutable y_pos: float;
               mutable x_vel: float;
               mutable y_vel: float;
               fb_radius: float;
               fb_mass: float;
             }

let cube x = x *. x *. x
let square x = x *. x
(* let sqrtcube x = sqrt (x *. x *. x)   (* Faster but less flexible
   implementaqtion of sqrtcube *) *)

let body_to_fbody body =
  let (x_pos,y_pos) = body.pos
  and (x_vel,y_vel) = body.velocity
  (*and (x_accel,y_accel) =
    match body.i with None -> (0.,0.) | Some accel -> accel *)
  in
  { x_pos = x_pos;
    y_pos = y_pos;
    fb_radius = body.radius;
    fb_mass = body.mass;
    x_vel = x_vel;
    y_vel = y_vel;
  }

let build_fbody_array bodies =
  Array.of_list (List.map ~f:body_to_fbody bodies)

(* convert a body and a b to an updated body *)
let fbody_and_body_to_body body b =
  { body with
    velocity = (b.x_vel,b.y_vel);
    pos = (b.x_pos,b.y_pos);
  }


(******** List iterator w/index ********)
let rec list_iteri_rec ~f list i = match list with
    [] -> ()
  | hd::tl -> f ~i hd; list_iteri_rec ~f tl (i+1)

let list_iteri ~f list =
  list_iteri_rec ~f list 0

(**************************************************)

let array_to_bodies bodies array =
  let b_list = Array.to_list array in
  List.map2 ~f:fbody_and_body_to_body bodies b_list

(**********************************************************************)

let timer = MTimer.create ()

let sqrtcube x = sqrt (x *. x *. x)

let act_all_on_all_rk ~bounce bodies =
  let gexp = grav_exp#v in
  let exp = ((1.0 +. gexp)/.2.0) in
  let sqrtcube =
    if gexp = 2.0 then sqrtcube else (fun x -> x ** exp) in
  let const = gconst#v in

  (* t is the time.  This function has no time dependence.
     s is the position is state space
     dsdt is the array of derivatives
  *)
  let deriv _t s dsdt =
    (* initialize derivatives to 0 *)
    for i = 0 to Array.length dsdt - 1 do dsdt.(i) <- 0. done;

    for i = 0 to Array.length bodies - 1 do
      (* x and y pos derivative *)
      dsdt.(i*4) <- s.(i*4 + 2);
      dsdt.(i*4 + 1) <- s.(i*4 + 3);
    done;

    for i = 0 to Array.length bodies - 1 do
      for j = i+1 to Array.length bodies - 1 do
        (* compute i's action on j and vice-versa.  That way you only need to
         * compute the force once.  It's nearly twice as fast that way. *)
        let x_pos_i = s.(i*4) and y_pos_i = s.(i*4 + 1) in
        let x_pos_j = s.(j*4) and y_pos_j = s.(j*4 + 1) in

        let xdiff = x_pos_i -. x_pos_j
        and ydiff = y_pos_i -. y_pos_j in

        let dist_sq = xdiff *. xdiff +. ydiff *. ydiff in
        let mult = const /. sqrtcube dist_sq in
        let mult_i = -. mult *. bodies.(j).fb_mass
        and mult_j = mult *. bodies.(i).fb_mass
        in

        (* x vel derivative *)
        dsdt.(i*4 + 2) <- dsdt.(i*4 + 2) +. xdiff *. mult_i;
        dsdt.(j*4 + 2) <- dsdt.(j*4 + 2) +. xdiff *. mult_j;

        (* y vel derivative *)
        dsdt.(i*4 + 3) <- dsdt.(i*4 + 3) +. ydiff *. mult_i;
        dsdt.(j*4 + 3) <- dsdt.(j*4 + 3) +. ydiff *. mult_j;


        if bounce &&
          (let total_radius = bodies.(j).fb_radius +. bodies.(i).fb_radius in
           dist_sq < total_radius *. total_radius )
        then (
          let total_radius = bodies.(j).fb_radius +. bodies.(i).fb_radius in

          let mult = -. mult *. bodies.(i).fb_mass *. bodies.(j).fb_mass
            *. (total_radius *. total_radius /. dist_sq) in
          let mult_i = -. mult /. bodies.(i).fb_mass
          and mult_j = mult /. bodies.(j).fb_mass
          in

          (* x vel derivative *)
          dsdt.(i*4 + 2) <- dsdt.(i*4 + 2) +. xdiff *. mult_i;
          dsdt.(j*4 + 2) <- dsdt.(j*4 + 2) +. xdiff *. mult_j;

          (* y vel derivative *)
          dsdt.(i*4 + 3) <- dsdt.(i*4 + 3) +. ydiff *. mult_i;
          dsdt.(j*4 + 3) <- dsdt.(j*4 + 3) +. ydiff *. mult_j;
        )

      done
    done
  in


  let s =
    Array.init (4 * Array.length bodies)
      ~f:(fun i -> match i mod 4 with
      | 0 -> bodies.(i / 4).x_pos
      | 1 -> bodies.(i / 4).y_pos
      | 2 -> bodies.(i / 4).x_vel
      | _ -> bodies.(i / 4).y_vel
      )
  in

  (* dumb_solver s state.delta#v deriv;  *)
  Rk4.step ~y:s ~x:0.0 ~h:state.delta#v ~yout:s ~derivs:deriv;

  for i = 0 to Array.length bodies - 1 do
    bodies.(i).x_pos <- s.(4 * i);
    bodies.(i).y_pos <- s.(4 * i + 1);
    bodies.(i).x_vel <- s.(4 * i + 2);
    bodies.(i).y_vel <- s.(4 * i + 3);
  done


let act_all_on_all ~bounce bodies =
  let array = build_fbody_array bodies in
  act_all_on_all_rk ~bounce array;
  array_to_bodies bodies array