Merge branch 'grid_shift'

pyfeyn2/auto/position.py

@@ -1,7 +1,239 @@
-from pyfeyn2.feynmandiagram import Propagator
+import itertools
+import logging
+from itertools import permutations
+
+import numpy as np
+from feynml import Point, Propagator
+
 from pyfeyn2.interface.dot import dot_to_positions, feynman_to_dot
 
 
+# from https://stackoverflow.com/a/9997374
+def ccw(A, B, C):
+    """
+    Return true if the points A, B, and C are in counter-clockwise order.
+    """
+    return (C.y - A.y) * (B.x - A.x) > (B.y - A.y) * (C.x - A.x)
+
+
+# Return true if line segments AB and CD intersect
+def intersect(A, B, C, D):
+    """
+    Return true if line segments AB and CD intersect
+
+    Parameters
+    ----------
+    A : Point
+        The first point of the first line segment.
+    B : Point
+        The second point of the first line segment.
+    C : Point
+        The first point of the second line segment.
+    D : Point
+        The second point of the second line segment.
+
+    Returns
+    -------
+    bool
+        True if the line segments intersect, False otherwise.
+
+    Examples
+    --------
+    >>> A = Point(0, 0)
+    >>> B = Point(1, 1)
+    >>> C = Point(0, 1)
+    >>> D = Point(1, 0)
+    >>> intersect(A, B, C, D)
+    True
+    >>> A,B,C,D = Point(0,0), Point(1,1), Point(0,0), Point(1,0)
+    >>> intersect(A, B, C, D)
+    False
+    """
+    if A.x == C.x and A.y == C.y:
+        return False
+    if A.x == D.x and A.y == D.y:
+        return False
+    if B.x == C.x and B.y == C.y:
+        return False
+    if B.x == D.x and B.y == D.y:
+        return False
+    return ccw(A, C, D) != ccw(B, C, D) and ccw(A, B, C) != ccw(A, B, D)
+
+
+def require_xy(points):
+    # check if a vertex or leg is missing a x or y position
+    for v in points:
+        if v.x is None:
+            raise Exception(f"Vertex or leg {v} is missing x position.")
+        if v.y is None:
+            raise Exception(f"Vertex or leg {v} is missing y position.")
+
+
+def _compute_number_of_intersects(fd):
+    """
+    Computes the number of crossed propagators/legs in a Feynman diagram
+    """
+    # check if a vertex or leg is missing a x or y position
+    points = [*fd.vertices, *fd.legs]
+    require_xy(points)
+    lines = []
+    for p in fd.propagators:
+        src = fd.get_point(p.source)
+        tar = fd.get_point(p.target)
+        lines.append([src, tar])
+    for l in fd.legs:
+        if l.is_incoming():
+            src = Point(l.x, l.y)
+            tar = fd.get_point(l.target)
+            lines.append([src, tar])
+        elif l.is_outgoing():
+            src = fd.get_point(l.target)
+            tar = Point(l.x, l.y)
+            lines.append([src, tar])
+
+    ci = 0
+    for i, l1 in enumerate(lines):
+        for _, l2 in enumerate(lines[i + 1 :]):
+            # test if the lines cross, without changing the lines
+            if intersect(l1[0], l1[1], l2[0], l2[1]):
+                ci += 1
+    return ci
+
+
+def auto_remove_intersections_by_align_legs(fd, adjust_points=False, size=5):
+    """
+    Automatically remove intersections by aligning the legs and reshufffling (permuting) them.
+    """
+    fd = auto_align_legs(fd)
+    if adjust_points:
+        fd = feynman_adjust_points(fd, size=size, clear_vertices=True)
+    min_intersections = np.inf
+    min_perm = 0
+    inc = [l for l in fd.legs if l.is_incoming()]
+    outc = [l for l in fd.legs if l.is_outgoing()]
+    xyin = [[l.x, l.y] for l in inc]
+    xyout = [[l.x, l.y] for l in outc]
+    # loop over all permutations of incoming and outgoing legs
+    for i, o in itertools.product(
+        set(permutations(range(len(inc)))), set(permutations(range(len(outc))))
+    ):
+        for xyi, l in zip(xyin, i):
+            inc[l].x = xyi[0]
+            inc[l].y = xyi[1]
+        for xyo, l in zip(xyout, o):
+            outc[l].x = xyo[0]
+            outc[l].y = xyo[1]
+        if adjust_points:
+            fd = feynman_adjust_points(fd, size=size, clear_vertices=True)
+        ci = _compute_number_of_intersects(fd)
+        print(ci)
+        logging.debug(f"auto_remove_intersections_by_align_legs: {ci}")
+        if ci < min_intersections:
+            min_intersections = ci
+            min_perm = (i, o)
+            logging.debug(f"auto_remove_intersections_by_align_legs: {ci}")
+            logging.debug(f"auto_remove_intersections_by_align_legs: {i} {o}")
+            logging.debug(f"auto_remove_intersections_by_align_legs: {xyin} {xyout}")
+    # use/return best permutation
+    for xyi, l in zip(xyin, min_perm[0]):
+        inc[l].x = xyi[0]
+        inc[l].y = xyi[1]
+    for xyo, l in zip(xyout, min_perm[1]):
+        outc[l].x = xyo[0]
+        outc[l].y = xyo[1]
+    if adjust_points:
+        fd = feynman_adjust_points(fd, size=size, clear_vertices=True)
+    return fd
+
+
+def auto_align_legs(fd, incoming=None, outgoing=None):
+    """
+    Automatically reshuffle the legs of a Feynman diagram.
+    """
+    f_min_x, f_min_y, f_max_x, f_max_y = fd.get_bounding_box()
+    inc = [l for l in fd.legs if l.is_incoming()]
+    outc = [l for l in fd.legs if l.is_outgoing()]
+    if incoming is None:
+        incoming = [[f_min_x, y] for y in np.linspace(f_min_y, f_max_y, len(inc))]
+    if outgoing is None:
+        outgoing = [[f_max_x, y] for y in np.linspace(f_min_y, f_max_y, len(outc))]
+    _auto_align(inc, incoming)
+    _auto_align(outc, outgoing)
+    return fd
+
+
+def _auto_align(points, positions):
+    """
+    Automatically position the vertices and legs on a list of positions.
+    """
+    logging.debug(f"_auto_align: positions {positions}")
+    # check if a vertex or leg is missing a x or y position
+    require_xy(points)
+    vpl = len(points)
+    # table of distances between vertices v and points p
+    dist = np.ones((vpl, len(positions))) * np.inf
+    for i, v in enumerate(points):
+        for j, p in enumerate(positions):
+            dist[i][j] = np.sqrt((v.x - p[0]) ** 2 + (v.y - p[1]) ** 2)
+    for i in range(vpl):
+        min_i, min_j = np.unravel_index(dist.argmin(), dist.shape)
+        v = points[min_i]
+        v.x = positions[min_j][0]
+        v.y = positions[min_j][1]
+        # remove min_i and min_j from dist
+        dist[min_i, :] = np.inf
+        dist[:, min_j] = np.inf
+
+
+def auto_align(fd, positions):
+    """
+    Automatically position the vertices and legs on a list of positions.
+
+    Parameters
+    ----------
+    fd : FeynmanDiagram
+        The Feynman diagram to be positioned.
+    positions : list of tuple
+        A list of tuples of the form (x,y) with the positions of the vertices
+
+    Returns
+    -------
+    FeynmanDiagram
+        The Feynman diagram with the vertices and legs positioned.
+    """
+    _auto_align([*fd.vertices, *fd.legs], positions)
+    return fd
+
+
+def auto_grid(fd, n_x=None, n_y=None, min_x=None, min_y=None, max_x=None, max_y=None):
+    """
+    Automatically position the vertices and legs on a grid, with the given
+    minimum and maximum values for x and y, and the number of grid points, but
+    avoid placing vertices or legs on the same position.
+    """
+    # get the bounding box and construct grid from that
+    f_min_x, f_min_y, f_max_x, f_max_y = fd.get_bounding_box()
+    if n_x is None:
+        n_x = len(fd.vertices) + len(fd.legs)
+    if n_y is None:
+        n_y = len(fd.vertices) + len(fd.legs)
+    if min_x is None:
+        min_x = f_min_x
+    if max_x is None:
+        max_x = f_max_x
+    if min_y is None:
+        min_y = f_min_y
+    if max_y is None:
+        max_y = f_max_y
+    logging.debug(f"auto_grid {n_x}, {n_y}, {min_x}, {min_y}, {max_x}, {max_y}")
+
+    xvalues = np.linspace(min_x, max_x, n_x)
+    yvalues = np.linspace(min_y, max_y, n_y)
+    xx, yy = np.meshgrid(xvalues, yvalues)
+    positions = [[x, y] for x, y in zip(xx.flatten(), yy.flatten())]
+    return auto_align(fd, positions)
+
+
 def auto_position(fd, layout="neato", clear_vertices=True):
     """Automatically position the vertices and legs."""
     # fd = scale_positions(fd, 10)

pyfeyn2/feynmandiagram.py

@@ -1,20 +1,19 @@
 """Moved to :py:mod:`feynml`"""
 from importlib.metadata import version
 
-from feynml.connector import Connector as Connector_
-from feynml.feynmandiagram import FeynmanDiagram as FeynmanDiagram_
-from feynml.feynml import FeynML as FeynML_
-from feynml.head import Head as Head_
-
 # from feynml.feynml import Tool as Tool_
-from feynml.leg import Leg as Leg_
-from feynml.meta import Meta as Meta_
-from feynml.momentum import Momentum as Momentum_
-from feynml.pdgid import PDG as PDG_
-from feynml.point import Point as Point_
-from feynml.propagator import Propagator as Propagator_
-from feynml.styled import Styled as Styled_
-from feynml.vertex import Vertex as Vertex_
+from feynml import PDG as PDG_
+from feynml import Connector as Connector_
+from feynml import FeynmanDiagram as FeynmanDiagram_
+from feynml import FeynML as FeynML_
+from feynml import Head as Head_
+from feynml import Leg as Leg_
+from feynml import Meta as Meta_
+from feynml import Point as Point_
+from feynml import Propagator as Propagator_
+from feynml import Styled as Styled_
+from feynml import Vertex as Vertex_
+from feynml.momentum import Momentum as Momentum_  # TODO fix to feynml only
 from smpl_doc import doc
 
 

pyproject.toml

@@ -29,7 +29,7 @@ cssselect ="*"
 smpl_io = "*"
 smpl_doc = "*"
 smpl_util= "*"
-feynml = {version = ">=0.1.6", extras = ["interfaces"]}
+feynml = {version = ">=0.2.8", extras = ["interfaces"]}
 #feynml= {path= "../feynml", develop = true, extras = ["interfaces"]}
 
 [tool.poetry.scripts]

tests/test_auto.py

@@ -0,0 +1,97 @@
+import numpy as np
+from feynml import FeynmanDiagram, Leg, Propagator, Vertex
+
+from pyfeyn2.auto.position import (
+    auto_align,
+    auto_align_legs,
+    auto_grid,
+    auto_remove_intersections_by_align_legs,
+    feynman_adjust_points,
+)
+
+
+def _get_fd_2_2():
+    v1 = Vertex("v1").with_shape("dot")
+    v2 = Vertex("v2").with_style("symbol : dot")
+
+    fd = FeynmanDiagram().add(
+        v1,
+        v2,
+        Propagator(name="g").connect(v1, v2),
+        Leg(name="g").with_target(v1).with_xy(-2, 1).with_incoming(),
+        Leg(name="g")
+        .with_target(v2)
+        .with_xy(-2, -1)
+        .with_incoming()
+        .with_class("notred"),
+        Leg(name="g").with_target(v1).with_xy(2, 1).with_outgoing().with_class("red"),
+        Leg("myid1", name="g").with_target(v2).with_xy(2, -1).with_outgoing(),
+    )
+    return fd
+
+
+def _get_fd_2_4():
+    v1 = Vertex("v1").with_shape("dot")
+    v2 = Vertex("v2").with_style("symbol : dot")
+
+    fd = FeynmanDiagram().with_rules(
+        """ * {color: red;} 
+            [type=fermion] {color: blue; line: gluon}
+            #p1 {color: green;}
+            :not([type=fermion]) { color : black; line: fermion}"""
+    )
+    v1 = Vertex("v1")
+    v2 = Vertex("v2")
+    v3 = Vertex("v3")
+    v4 = Vertex("v4")
+    p1 = Propagator("p1").connect(v1, v2).with_type("gluon")
+    p2 = Propagator("p2").connect(v1, v3).with_type("gluon")
+    p3 = Propagator("p3").connect(v3, v2).with_type("gluon")
+    p4 = Propagator("p4").connect(v4, v3).with_type("gluon")
+    p5 = Propagator("p5").connect(v4, v2).with_type("gluon")
+    l1 = Leg("l1").with_target(v1).with_type("gluon").with_incoming().with_xy(2, 1)
+    l2 = Leg("l2").with_target(v1).with_type("gluon").with_incoming().with_xy(-2, -1)
+    l3 = (
+        Leg("l3")
+        .with_target(v2)
+        .with_type("fermion")
+        .with_outgoing()
+        .with_xy(2, -2)
+        .with_class("blue")
+    )
+    l4 = Leg("l4").with_target(v3).with_type("fermion").with_outgoing().with_xy(2, 2)
+    l5 = Leg("l5").with_target(v4).with_type("gluon").with_outgoing().with_xy(2, 1)
+    l6 = Leg("l6").with_target(v4).with_type("gluon").with_outgoing().with_xy(-2, -1)
+
+    l6.style.color = "orange"
+
+    fd.propagators.extend([p1, p2, p3, p4, p5])
+    fd.vertices.extend([v1, v2, v3, v4])
+    fd.legs.extend([l1, l2, l3, l4, l5, l6])
+    return fd
+
+
+def test_auto_grid():
+    fd = _get_fd_2_2()
+    fd = feynman_adjust_points(fd, size=10)
+    fd = auto_grid(fd, n_x=3, n_y=4)
+
+
+def test_auto_align():
+    fd = _get_fd_2_2()
+    fd = feynman_adjust_points(fd, size=10)
+    fd = auto_align(
+        fd, np.array([[-1, -1], [1, -1], [1, 1], [-1, 1], [0, -0.05], [0.0, 0.05]]) * 2
+    )
+
+
+def test_auto_align_legs():
+    fd = _get_fd_2_4()
+    fd = feynman_adjust_points(fd, size=10)
+    fd = auto_align_legs(fd)
+
+
+def test_auto_remove_intersections_by_align_legs():
+    fd = _get_fd_2_4()
+    fd = feynman_adjust_points(fd, size=10)
+    fd = auto_remove_intersections_by_align_legs(fd)