QCaml/Utils/Coordinate.ml

open Util

type t = 
  | Bohr     of (float * float * float)
  | Angstrom of (float * float * float)

let zero = Bohr (0., 0., 0.)

let of_float_triplet (x,y,z) = function
| `Bohr     -> Bohr     (x,y,z)
| `Angstrom -> Angstrom (x,y,z)

let of_3_floats x y z =
  of_float_triplet (x,y,z)

let to_string t =
  let result (x,y,z) = 
    (string_of_float x)^" "^(string_of_float y)^" "^(string_of_float z)
  in
  match t with
  | Bohr     x -> (result x) ^ " Bohr"
  | Angstrom x -> (result x) ^ " Angstrom"


let extract_float_tuple = function
  | Bohr     a
  | Angstrom a -> a


(** Linear algebra *)
let (|.) s a =
  match a with
  | Bohr     (x,y,z) -> Bohr     ( s*.x, s*.y, s*.z )
  | Angstrom (x,y,z) -> Angstrom ( s*.x, s*.y, s*.z )

let to_Angstrom = function
  | Angstrom a -> Angstrom a
  | Bohr     a -> Angstrom (Constants.a0 |. Bohr a |> extract_float_tuple)

let to_Bohr = function
  | Angstrom a -> Bohr (1./.Constants.a0 |. Angstrom a |> extract_float_tuple)
  | Bohr     a -> Bohr a 

let (|-), (|+) =
  let rec op f p q =
    match (p, q) with
    | (Angstrom a, Angstrom b) -> Angstrom (f a b)
    | (Bohr     a, Bohr     b) -> Bohr     (f a b)
    | (Angstrom a, Bohr     b) -> op f (to_Bohr p) q
    | (Bohr     a, Angstrom b) -> op f p (to_Bohr q)
  in
  (op (fun (x,y,z) (x',y',z') -> ( x-.x', y-.y', z-.z' )) , 
   op (fun (x,y,z) (x',y',z') -> ( x+.x', y+.y', z+.z' ))
  )


let neg a = -1. |. a

let rec dot p q =
  match (p,q) with 
    | Bohr (x,y,z), Bohr (x',y',z') ->  x*.x' +. y*.y' +. z*.z' 
    | _ -> dot (to_Bohr p) (to_Bohr q)


let norm u =
  sqrt @@ dot u u


let rec to_tuple a = 
  to_Bohr a |> extract_float_tuple

let x a = 
  let (result, _, _) = extract_float_tuple @@ to_Bohr a in
  result

let y a = 
  let (_, result, _) = extract_float_tuple @@ to_Bohr a in
  result

let z a = 
  let (_, _, result) = extract_float_tuple @@ to_Bohr a in
  result

let coord a = function
  | 0 -> x a
  | 1 -> y a
  | 2 -> z a 
  | _ -> raise (Invalid_argument "Coordinate")