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
  )

Here is a quick program I wrote to enumerate all Feynman diagrams within the rules set by this PBS SpaceTime episode, to the best of my understanding of those rules.

You will need an up to date OCaml compiler with the Core library to run it. You can find instructions here if you're new to the language. (It's one of the best computer languages. If you are interested in learning it, you will find an entire online book at that address.)

This program was used to list all diagrams in the picture below except the first two, after which I used JaxoDraw to actually draw them. The ones in the bottom left corner were not included in the official answer, maybe because they are considered self-energetic interactions?