module StringMap = Map.Make(String) type atom_id = int type angle = Label of string | Value of float type distance = Label of string | Value of float type dihedral = Label of string | Value of float let pi = acos (-1.) let to_radian = pi /. 180. let rec in_range (xmin, xmax) x = if (x <= xmin) then in_range (xmin, xmax) (x -. xmin +. xmax ) else if (x > xmax) then in_range (xmin, xmax) (x -. xmax +. xmin ) else x let atom_id_of_int : int -> atom_id = fun x -> ( assert (x>0) ; x) let distance_of_float : float -> distance = fun x -> ( assert (x>=0.) ; Value x) let angle_of_float : float -> angle = fun x -> Value (in_range (-180., 180.) x) let dihedral_of_float : float -> dihedral = fun x -> Value (in_range (-360., 360.) x) let atom_id_of_string : string -> atom_id = fun i -> atom_id_of_int @@ int_of_string i let distance_of_string : string -> distance = fun s -> try distance_of_float @@ float_of_string s with _ -> Label s let angle_of_string : string -> angle = fun s -> try angle_of_float @@ float_of_string s with _ -> Label s let dihedral_of_string : string -> dihedral = fun s -> try dihedral_of_float @@ float_of_string s with _ -> Label s let int_of_atom_id : atom_id -> int = fun x -> x let float_of_distance : float StringMap.t -> distance -> float = fun map -> function | Value x -> x | Label s -> StringMap.find s map let float_of_angle : float StringMap.t -> angle -> float = fun map -> function | Value x -> x | Label s -> StringMap.find s map let float_of_dihedral : float StringMap.t -> dihedral -> float = fun map -> function | Value x -> x | Label s -> StringMap.find s map type line = | First of Element.t | Second of (Element.t * distance) | Third of (Element.t * atom_id * distance * atom_id * angle) | Other of (Element.t * atom_id * distance * atom_id * angle * atom_id * dihedral ) | Coord of (string * float) let string_of_line map = let f_r = float_of_distance map and f_a = float_of_angle map and f_d = float_of_dihedral map and i_i = int_of_atom_id in function | First e -> Printf.sprintf "%-3s" (Element.to_string e) | Second (e, r) -> Printf.sprintf "%-3s %5d %f" (Element.to_string e) 1 (f_r r) | Third (e, i, r, j, a) -> Printf.sprintf "%-3s %5d %f %5d %f" (Element.to_string e) (i_i i) (f_r r) (i_i j) (f_a a) | Other (e, i, r, j, a, k, d) -> Printf.sprintf "%-3s %5d %f %5d %f %5d %f" (Element.to_string e) (i_i i) (f_r r) (i_i j) (f_a a) (i_i k) (f_d d) | Coord (c, f) -> Printf.sprintf "%s %f" c f let line_of_string l = let line_clean = Str.split (Str.regexp " ") l |> List.filter (fun x -> x <> "") in match line_clean with | e :: [] -> First (Element.of_string e) | e :: i :: r :: [] -> Second (Element.of_string e, distance_of_string r) | e :: i :: r :: j :: a :: [] -> Third (Element.of_string e, atom_id_of_string i, distance_of_string r, atom_id_of_string j, angle_of_string a) | e :: i :: r :: j :: a :: k :: d :: [] -> Other (Element.of_string e, atom_id_of_string i, distance_of_string r, atom_id_of_string j, angle_of_string a, atom_id_of_string k, dihedral_of_string d) | c :: f :: [] -> Coord (c, float_of_string f) | _ -> failwith ("Syntax error: "^l) type t = (line array * float StringMap.t) let of_string t = let l = Str.split (Str.regexp "\n") t |> List.map String.trim |> List.filter (fun x -> x <> "") |> List.map line_of_string in let l = match l with | First _ :: Second _ :: Third _ :: _ | First _ :: Second _ :: Coord _ :: [] | First _ :: Second _ :: [] | First _ :: [] -> l | _ -> failwith "Syntax error" in let (l, m) = let rec work lst map = function | (First _ as x) :: rest | (Second _ as x) :: rest | (Third _ as x) :: rest | (Other _ as x) :: rest -> work (x::lst) map rest | (Coord (c,f)) :: rest -> work lst (StringMap.add c f map) rest | [] -> (List.rev lst, map) in work [] (StringMap.empty) l in (Array.of_list l, m) (** Linear algebra *) let (|-) (x,y,z) (x',y',z') = ( x-.x', y-.y', z-.z' ) let (|+) (x,y,z) (x',y',z') = ( x+.x', y+.y', z+.z' ) let (|.) s (x,y,z) = ( s*.x, s*.y, s*.z ) let dot (x,y,z) (x',y',z') = x*.x' +. y*.y' +. z*.z' let norm u = sqrt @@ dot u u let normalized u = 1. /. (norm u) |. u let cross (x,y,z) (x',y',z') = ((y *. z' -. z *. y'), -. (x *. z' -. z *. x'), (x *. y' -. y *. x')) let rotation_matrix axis angle = (* Euler-Rodrigues formula for rotation matrix, taken from https://github.com/jevandezande/zmatrix/blob/master/converter.py *) let a = (cos (angle *. to_radian *. 0.5)) in let (b, c, d) = (-. sin (angle *. to_radian *. 0.5)) |. (normalized axis) in Array.of_list @@ [(a *. a +. b *. b -. c *. c -. d *. d, 2. *. (b *. c -. a *. d), 2. *. (b *. d +. a *. c)); (2. *. (b *. c +. a *. d), a *. a +. c *. c -.b *. b -. d *. d, 2. *. (c *. d -. a *. b)); (2. *. (b *. d -. a *. c), 2. *. (c *. d +. a *. b), a *. a +. d *. d -. b *. b -. c *. c)] let apply_rotation_matrix rot u = (dot rot.(0) u, dot rot.(1) u, dot rot.(2) u) let to_xyz (z,map) = let result = Array.make (Array.length z) None in let get_cartesian_coord i = match result.(i-1) with | None -> failwith @@ Printf.sprintf "Atom %d is defined in the future" i | Some (_, x, y, z) -> (x, y, z) in let append_line i' = match z.(i') with | First e -> result.(i') <- Some (e, 0., 0., 0.) | Second (e, r) -> let r = float_of_distance map r in result.(i') <- Some (e, 0., 0., r) | Third (e, i, r, j, a) -> begin let i, r, j, a = int_of_atom_id i, float_of_distance map r, int_of_atom_id j, float_of_angle map a in let ui, uj = get_cartesian_coord i, get_cartesian_coord j in let u_ij = (uj |- ui) in let rot = rotation_matrix (0., 1., 0.) a in let new_vec = apply_rotation_matrix rot ( r |. (normalized u_ij)) in let (x, y, z) = new_vec |+ ui in result.(i') <- Some (e, x, y, z) end | Other (e, i, r, j, a, k, d) -> begin let i, r, j, a, k, d = int_of_atom_id i, float_of_distance map r, int_of_atom_id j, float_of_angle map a, int_of_atom_id k, float_of_dihedral map d in let ui, uj, uk = get_cartesian_coord i, get_cartesian_coord j, get_cartesian_coord k in let u_ij, u_kj = (uj |- ui) , (uj |- uk) in let normal = cross u_ij u_kj in let new_vec = r |. (normalized u_ij) |> apply_rotation_matrix (rotation_matrix normal a) |> apply_rotation_matrix (rotation_matrix u_ij d) in let (x, y, z) = new_vec |+ ui in result.(i') <- Some (e, x, y, z) end | Coord _ -> () in Array.iteri (fun i _ -> append_line i) z; let result = Array.map (function | Some x -> x | None -> failwith "Some atoms were not defined" ) result in result let to_xyz_string (l,map) = String.concat "\n" ( to_xyz (l,map) |> Array.map (fun (e,x,y,z) -> Printf.sprintf "%s %f %f %f\n" (Element.to_string e) x y z) |> Array.to_list )