PBS SpaceTime Feynman Diagram Challenge

feynman.ml

(*
Compute all possible Feynman diagrams with 1 el + 1 pos input,
1 el + 1 pos output, 4 nodes, and no self-energetic transitions.
*)
open Core.Std

type particle = Electron | Positron | Photon [@@deriving compare, sexp]

let inverse = function
  | Photon   -> Photon
  | Electron -> Positron
  | Positron -> Electron

type node = int [@@deriving compare, sexp]

type target = Input | Output | Node of node [@@deriving compare, sexp]

(* a Feynman diagram is a map (node * particle -> target) *)
module Feynman = struct
  module K = struct
    module T = struct
      type t = node * particle [@@deriving compare, sexp]
    end
    include T
    include Comparable.Make (T)
  end

  module T = struct
    type t = target K.Map.t [@@deriving compare, sexp]
  end
  include K.Map
  include Comparable.Make (T)
end

(* utility Set of (node * node) *)
module Node2 = struct
  module T = struct
    type t = node * node [@@deriving compare, sexp]
  end
  include T
  include Comparable.Make (T)
end

let n_nodes = 4

(* check a Feynman diagram that is being built (by adding edges) and return:
   `Ok      if it this is a valid Feynman diagram
   `Not_yet if it's incomplete, but adding the right edges may make it right
   `Stop    if it's not consistent: adding more edges will never make it right
*)
let check graph =
  let count ~f = Map.counti graph ~f:(fun ~key:(_, p) ~data:t -> f t p) in
  let node2set ~f = Map.fold graph ~init:Node2.Set.empty
    ~f:(fun ~key:(n, p) ~data:t set ->
      match t with
      | Node n' when f p -> Set.add set (n, n')
      | _ -> set
    ) in
  let i_ph  = count ~f:(fun t p -> t = Input  && p = Photon  ) in
  let i_el  = count ~f:(fun t p -> t = Input  && p = Electron) in
  let i_po  = count ~f:(fun t p -> t = Input  && p = Positron) in
  let o_ph  = count ~f:(fun t p -> t = Output && p = Photon  ) in
  let o_el  = count ~f:(fun t p -> t = Output && p = Electron) in
  let o_po  = count ~f:(fun t p -> t = Output && p = Positron) in
  let ph    = node2set ~f:(fun p -> p = Photon) in
  let el_po = node2set ~f:(fun p -> p <> Photon)
  in
  if i_ph > 0 || i_el > 1 || i_po > 1 || o_ph > 0 || o_el > 1 || o_po > 1
     || not (Node2.Set.is_empty (Node2.Set.inter el_po ph))
  then `Stop
  else if i_el < 1 || i_po < 1 || o_el < 1 || o_po < 1
  then `Not_yet
  else `Ok

(* Poor Man's textual representation of a Feynman diagram *)
let graph_to_string graph =
  Map.to_alist graph
  |> List.filter ~f:(fun ((n, _), t) ->
    match t with
      | Node n' -> n < n'
      | _ -> true
  )
  |> List.map ~f:(fun ((n, p), t) ->
    let p_lab = match p with
      | Photon   -> "~"
      | Electron -> ">"
      | Positron -> "<"
    in let t_lab = match t with
      | Input  -> "I"
      | Output -> "O"
      | Node n -> string_of_int n
    in
    string_of_int n ^ p_lab ^ t_lab
  )
  |> String.concat ~sep:", "

(* given a Feynman diagram and a permutation of the numbers 1..4, return the
   same graph with the node numbers changed according to the permutation *)
let graph_perm graph perm =
  Map.fold graph ~init:Feynman.empty ~f:(fun ~key:(n, p) ~data:t new_graph ->
    let n' = List.nth_exn perm (n - 1) in
    let t' = match t with
      | Node n -> Node (List.nth_exn perm (n - 1))
      | c -> c
    in
    Map.add new_graph ~key:(n', p) ~data:t'
  )

(* call ~f on all possible graphs that satisfy the rules, including both
   implicit rules and those implemented in check *)
let iter_all_graphs ~f =
  (* start with the given graph and try every possible connection starting from
     the given node and particle, using only subsequent nodes (or input/output)
     as targets, but not previous nodes *)
  let rec walk node particle graph =
    let key = (node, particle) in
    let next = match particle with
      | Photon   -> walk node Electron
      | Electron -> walk node Positron
      | Positron -> if node < n_nodes
          then walk (node + 1) Photon
          else fun graph -> if check graph = `Ok then f graph
    in
    if check graph = `Stop then ()
    else if Map.mem graph key then
      next graph
    else begin
      Map.add graph ~key ~data:Input |> next;
      Map.add graph ~key ~data:Output |> next;
      for n = node + 1 to n_nodes do
        let inv_key = (n, inverse particle)
        in
        if not (Map.mem graph inv_key) then
          Map.add graph ~key ~data:(Node n)
          |> Map.add ~key:inv_key ~data:(Node node)
          |> next
      done
    end
  in
  walk 1 Photon Feynman.empty

(* Heap's algorithm: generate a list of all permutations of a list of items *)
let permutations lst =
  let out = ref [] in
  let a = Array.of_list lst in
  let rec generate n =
    if n = 1 then
      out := Array.to_list a :: !out
    else begin
      for i = 0 to n - 2 do
        generate (n - 1);
        if n % 2 = 0
        then Array.swap a i (n-1)
        else Array.swap a 0 (n-1)
      done;
      generate (n - 1)
    end
  in
  generate (Array.length a);
  !out

let all_perm = permutations [1;2;3;4]

let () =
  let seen = ref Feynman.Set.empty
  in
  iter_all_graphs ~f:(fun graph ->
    if not (Set.mem !seen graph)
    then begin
      seen := Feynman.Set.union !seen (
        List.map all_perm ~f:(graph_perm graph)
        |> Feynman.Set.of_list
      );
      printf "%s\n" (graph_to_string graph)
    end
  )