QCaml/Utils/Coordinate.ml

98 lines
2.3 KiB
OCaml

type t =
| Bohr of float array
| Angstrom of float array
(** Bohr radius, taken from https://physics.nist.gov/cgi-bin/cuu/Value?bohrrada0 *)
let a0 =
0.529_177_210_67
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 y =
let result x =
(string_of_float x.(0))^" "^(string_of_float x.(1))^" "^(string_of_float x.(2))
in
match y with
| Bohr x -> (result x) ^ " Bohr"
| Angstrom x -> (result x) ^ " Angstrom"
let to_float_array = function
| Bohr a
| Angstrom a -> a
let x a = (to_float_array a).(0)
let y a = (to_float_array a).(1)
let z a = (to_float_array a).(2)
let coord a = function
| 0 -> (to_float_array a).(0)
| 1 -> (to_float_array a).(1)
| 2 -> (to_float_array a).(2)
| _ -> raise (Invalid_argument "Coordinate")
(** 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 |]
| _ -> assert false
let to_Angstrom = function
| Angstrom a -> Angstrom a
| Bohr a -> Angstrom (a0 |. Bohr a |> to_float_array)
let to_Bohr = function
| Angstrom a -> Bohr (1./.a0 |. Angstrom a |> to_float_array)
| 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 a b ->
match a,b with
| [|x;y;z|], [|x';y';z'|] -> [| x-.x'; y-.y'; z-.z' |]
| _ -> assert false
) ,
op (fun a b ->
match a,b with
| [|x;y;z|], [|x';y';z'|] -> [| x+.x'; y+.y'; z+.z' |]
| _ -> assert false
)
)
let dot =
let rec op f p q =
match (p, q) with
| (Angstrom a, Angstrom b)
| (Bohr a, Bohr b) -> 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 a b ->
match a,b with
| [|x;y;z|], [|x';y';z'|] -> x*.x' +. y*.y' +. z*.z'
| _ -> assert false
)
let norm u =
sqrt @@ dot u u