laserprojector animation

projectoranim.py

import math
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d as a3
from matplotlib.animation import FuncAnimation
from dataclasses import dataclass
from collections import defaultdict
import numpy

"""
Author: Willem Hengeveld <[email protected]>

Drawing several objects:
    - prisms in multiple orientations and sizes.
    - lines for the beam
    - calculates two reflections, and finally the projection onto a screen.
    - while and animating the rotating prisms.
"""

class Point:
    """
    A single point.
    """
    def __init__(self, *args, **kwargs):
        self.color = kwargs.get('color', [0,0,0,1])
        if len(args)==0:
            self.coord = []
        elif len(args)==1:
            arg = args[0]
            try:
                self.coord = list(arg)
            except TypeError:
                # handle 1-d Point(1) or Point(1.1)
                self.coord = [arg]
        else:
            self.coord = list(args)
    def __add__(lhs, rhs):
        return Point(l+r for l, r in zip(lhs.coord, rhs.coord))
    def __sub__(lhs, rhs):
        return Point(l-r for l, r in zip(lhs.coord, rhs.coord))
    def __rmul__(rhs, lhs):
        return Point(lhs*x for x in rhs.coord)
    def __mul__(lhs, rhs):
        return Point(x*rhs for x in lhs.coord)
    def __truediv__(lhs, rhs):
        return Point(x/rhs for x in lhs.coord)
    def __floordiv__(lhs, rhs):
        return Point(x//rhs for x in lhs.coord)
    def __eq__(lhs, rhs):
        return lhs.dim()==rhs.dim() and all(a==b for a, b in zip(lhs.coord, rhs.coord))
    def __ne__(lhs, rhs):
        return not (lhs==rhs)
    def dim(self):
        return len(self.coord)

    def __getitem__(self, i):
        return self.coord[i]
    def __setitem__(self, i, x):
        self.coord[i] = x

    def __hash__(self):
        return hash(tuple(self.coord))

    @staticmethod
    def frompolar(angle, u, v):
        """
        u, v are perpendicular base vectors.
        """
        return math.cos(angle)*u + math.sin(angle)*v
    def __repr__(self):
        return "("+",".join(map(str, self.coord))+")"
    def length(self):
        return math.sqrt(sum(_**2 for _ in self.coord))
    def np(self):
        return numpy.array(self.coord)
    def draw(self, ax):
        ax.scatter([self[0]], [self[1]], color=self.color)


def innerproduct(a, b):
    return sum(a*b for a, b in zip(a.coord, b.coord))


def perpendicularBase(Q):
    """
    calculate an arbitrary set of vectors perpendicular to the given 'Q'

    see https://stackoverflow.com/questions/36760771/how-to-calculate-a-circle-perpendicular-to-a-vector
    """
    L = Q.length()
    x, y, z = Q.coord
    sigma = L if x>0.0 else -L
    h = x + sigma
    beta = -1.0/(sigma*h)

    f = beta*y
    yield Point(f*h, 1.0+f*y, f*z)

    g = beta*z
    yield Point(g*h, g*y, 1.0+g*z)


@dataclass
class LineSegment:
    """
    points on the linesegment:
    all points: { p1*(1-a)+p2*a, a in [0..1] }

    P = p1*(1-a)+p2*a  = p1 -a*p1+a*p2 = p1 + a*(p2-p1)
    --> P-p1 = a*(p2-p1)

    in 2d, and eliminating 'a' you get:
    
    (P.y-p1.y)*(p2.x-p1.x)= (P.x-p1.x)*(p2.y-p1.y)
    or 
    P.x*(p1.y-p2.y) + P.y*(p2.x-p1.x) = p1.y*p2.x -p1.x*p2.y 


    intersection P of two lines l and m:
       where l is the line through l1 and l2
       where m is the line through m1 and m2

    P = a*(l2-l1) + l1 = b*(m2-m1) + m1

    solve a, b from : a*(l2-l1) + l1 = b*(m2-m1) + m1
       a*(l2-l1) + b*(m1-m2) =  m1 - l1

       matrix equation:

       [  (l2-l1),  (m1-m2) ] * {a,b} = [ (m1-l1) ]

    """
    p0 : Point
    p1 : Point

    def draw(self, ax):
        ax.plot(*list( (x0, x1) for x0, x1 in zip(self.p0.coord, self.p1.coord)), color=[0,0,0,1])

    def move(self, delta):
        self.p0 += delta
        self.p1 += delta
    def pointForParams(self, a):
        return self.p0 + (self.p1-self.p0)*a


def isintriangle(p, a, b, c):
    """
    checks if the point 'p' is inside the triangle formed by {a,b,c}.

    all points must be 2D.
    """
    def sign(p1, p2, p3):
        return (p1[0]-p3[0])*(p2[1]-p3[1]) - (p2[0]-p3[0])*(p1[1]-p3[1])
    d1 = sign(p, a, b)
    d2 = sign(p, b, c)
    d3 = sign(p, c, a)
    has_neg = (d1<0) or (d2<0) or (d3<0)
    has_pos = (d1>0) or (d2>0) or (d3>0)

    return not (has_neg and has_pos)


class Polygon:
    """
    a collection of points which are all in the same plane.

    the plane is defined by three points: p1,p2,p3

    all points P = p1 + (p2-p1)*a + (p3-p1)*b
          for all a, b in R*R
        

    in 3d: the normal vector would be:
        N = (p2-p1) x (p3-p1)   -- cross product
    then the plane can be written as:

        all points P, which satisfy:
           N . (P-p1) == 0         -- dot product

    """
    def __init__(self, *args, **kwargs):
        self.color = kwargs.get('color', [0,0,0,1])
        if len(args)==0:
            self.points = []
        elif len(args)==1:
            arg = args[0]
            self.points = list(arg)
        else:
            self.points = list(args)
            
    def add(self, pt):
        self.points.append(pt)

    def draw(self, ax):
        vtx = self.vertices()
        tri = a3.art3d.Poly3DCollection([vtx])
        tri.set_color(self.color)
        ax.add_collection3d(tri)

    def vertices(self):
        return list(_.coord for _ in self.points)
    def __repr__(self):
        return "["+",".join(map(repr, self.points))+"]"

    def normalvector(self):
        """
        return a numpy unit normal vector for the plane.
        """
        p0, p1, p2 = self.points[:3]
        v = numpy.cross((p1-p0).np(), (p2-p0).np())
        return v / numpy.linalg.norm(v)

    def intersection(self, line):
        """
        calc the intersection point of the plane of the poly with the line.

        the intersection of a line l with this plane:
            line:  P = a*(l2-l1) + l1
            plane: N . (P-p1) == 0

            N . (a*(l2-l1) + l1-p1) == 0

            a == (N . (p1-l1)) / (N . (l2-l1))
        """
        N = Point(self.normalvector())
        nl = innerproduct(N, line.p1-line.p0)
        if nl==0:
            # no intersection
            return

        np = innerproduct(N, self.points[0]-line.p0)
        a = np/nl

        return line.pointForParams(a)

    def contains(self, point):
        """
        Check if point is within the boundaries of the polygon
        """
        n = Point(self.normalvector())

        # determine which axis is NOT paralel to the plane of the polygon.
        npcoord = None
        for c in range(point.dim()):
            axis = Point(int(c==i) for i in range(point.dim()))
            if innerproduct(axis, n) > 0.00001:
                npcoord = c
                break
        coordlist = set(range(point.dim())) - {npcoord}

        # reduce to 2D by removing that coordinate.
        i, j = list(coordlist)[:2]

        p = Point(point[i], point[j])
        a = b = c = None

        # enumerate all triangles fanning out from the first point.
        for q in self.points:
            if a is None:
                a = Point(q[i], q[j])
            else:
                b, c = c, Point(q[i], q[j])
            if b is not None:
                if isintriangle(p, a, b, c):
                    return True

class Shape:
    """
    Baseclass for a shape.
    a shape is a collection of polygons.

    The shape needs to implement two methods:
      - generatePoints, which gets passed all non-keyword arguments from the constructor.
      - generateFaces, which should generate the list of polygons which for the faces of the shape.
    """
    def __init__(self, *args, **kwargs):
        self.color = kwargs.get('color', [0,0,0,1])
        self.points = list(self.generatePoints(*args))
        self.faces = list(self.generateFaces())

    def draw(self, ax):
        for f in self.faces:
            f.draw(ax)

    def __repr__(self):
        return "<"+"\n ".join(map(repr, self.points))+">" # "<"+"\n ".join(f"{id(_)}" for _ in self.faces)+">"


class Tetrahedron(Shape):
    """
    generate all points for a tetrahedron
    """
    def generatePoints(self):
        yield Point(0, 0, 0)
        for i in range(3):
            yield Point(int(i!=j) for j in range(3))

    def generateFaces(self):
        for p in self.points:
            yield Polygon(set(self.points) - {p}, color=self.color.copy())


class Prism(Shape):
    """
    Generate a prism, with the specified axis, starting angle, nr of sides and radius.
    """
    def generatePoints(self, axis, startangle, n, radius):
        u, v = perpendicularBase(axis.p1-axis.p0)
        for i in range(n):
            p = Point.frompolar(startangle + i * 2*math.pi/n, radius*u, radius*v)
            yield axis.p0 + p
            yield axis.p1 + p 

    def generateFaces(self):
        yield Polygon(self.points[::2], color=self.color.copy())    # bottom
        yield Polygon(self.points[1::2], color=self.color.copy())   # top

        plist = self.points * 2

        for i in range(len(self.points)//2):
            corners = plist[2*i:2*i+4]
            corners[2:] = corners[4:1:-1]
            yield Polygon(corners, color=self.color.copy())

    def intersection(self, line):
        """
        # for all faces find the intersection point.
        # filter out faces where the intersection point is outside of the 
        #  boundaries of the face.
        # then find the one closest to the starting point of the line.

        return a face and a point.
        """
        found = []
        for f in self.faces:
            p = f.intersection(line)
            if p and f.contains(p):
                #print("found", p, " in ", f, " distance=", (p-line.p0).length(), (p-line.p1).length())
                found.append((f, p))
            elif p:
                #print("not in poly:", p, " in ", f, " distance=", (p-line.p0).length(), (p-line.p1).length())
                pass

        if not found:
            return None, None
        return min(found, key=lambda fp:(fp[1]-line.p0).length())


def calc_reflection(poly, incomingline):
    """
    Calculate a reflection using a householder transformation.

    note: bug, sometimes the reflection returned passes through the back of the mirror.

    return a LineSegment
    """
    intersectionpoint = poly.intersection(incomingline)

    normalvector = poly.normalvector()

    HH = numpy.identity(3) - 2*numpy.outer(normalvector, normalvector)
    pnew = numpy.matmul(HH, intersectionpoint.np()-incomingline.p0.np()) + intersectionpoint.np()

    return LineSegment(intersectionpoint, Point(pnew))
    

class Path:
    """
    Keep track of a set of points.
    """
    def __init__(self):
        self.path = []
        self.color = [0,0,0,1]
    def append(self, p):
        self.path.append(p)
    def draw(self, ax):
        ax.scatter(list(p[0] for p in self.path), list(p[1] for p in self.path), list(p[2] for p in self.path), color=self.color)


class World:
    """
    Keep track of all objects.
    """
    def __init__(self):

        # axis of the primary mirror, pointing up from the origin
        self.a = LineSegment(Point(0,0,2), Point(0,0,0))
        self.acolor = [1, 0, 0, 0.2]  # transparent red

        # axis of the secondary mirror, located 6 away on the y-axis, halfway the primary prism.
        self.b = LineSegment(Point(4,6,1), Point(-4,6,1))
        self.bcolor = [0, 1, 0, 0.4]  # transparent green

        # starting from the side of the small prism
        self.c = LineSegment(Point(3, 6, 1), Point(-1, 0, 1))
        self.ccolor = [0, 0, 1, 0.2]  # transparent blue
        self.path = Path()

    def update(self, angle):
        """
        update world with the given rotation angle.
        
        """

        # create the mirrors, 'a' rotates approx 40 times faster than 'b'.
        a = Prism(self.a, 4*angle/180*math.pi, 4, 4.0, color=self.acolor.copy())
        b = Prism(self.b, -0.1*angle/180*math.pi, 6, 2.0, color=self.bcolor.copy())

        # create the laser beam
        l0 = self.c
        l0.color = self.ccolor

        # create the view screen
        s = Polygon(Point(-10, -10, -10), Point(-10, -10, 10), Point(10, -10, 10), Point(10, -10, -10), color=[1,1,1,0.5])

        self.objects = [ a, b, s ]

        self.calculate_reflection(s, a, b, l0)


    def calculate_reflection(self, screen, a, b, l0):
        # calc intersection-face F0 of 'l0' with prism 'a'
        # calc intersection-point P0 of 'l0' with F0
        face0, point0 = a.intersection(l0)
        if not face0:
            print("no face0")
            return
        face0.color[3] = 0.8
        self.objects.append(LineSegment(l0.p0, point0))

        # calc reflection like l1 at P0
        l1 = calc_reflection(face0, l0)

        # calc intersection-face F1 of 'l1' with prism 'b'
        # calc intersection-point P1 of 'l1' with F1
        face1, point1 = b.intersection(l1)
        if not face1:
            #print("no face1")
            return
        face1.color[3] = 0.8
        self.objects.append(LineSegment(point0, point1))

        # calc reflection like l2 at P1
        l2 = calc_reflection(face1, l1)

        # calc intersection-point point2 of l2 with the screen
        point2 = screen.intersection(l2)
        if not point2:
            print("no point2")
            return
        self.objects.append(LineSegment(point1, point2))

        # remember all points on the screen.
        self.path.append(point2)
        self.objects.append(self.path)

    def draw(self, ax):
        for o in self.objects:
            o.draw(ax)

def usage():
    print("""Usage:
Use 'a' to select the primary/red mirror.
Use 'b' to select the secondary/green mirror.
Use 'c' to select the laser position.
Use ESC to deselect.
Then use cursor keys to move the position left/right, up/down.
""")

def main():
    usage()

    fig = plt.figure(figsize=(10,10))
    ax = fig.add_subplot(projection='3d')

    world = World()

    selected = None

    def on_press(e):
        nonlocal selected
        if e.key == 'ctrl+q':
             plt.close()
        elif e.key == 'a':
            world.acolor = [1, 0, 0, 0.5]  # transparent red
            world.bcolor = [0, 1, 0, 0.4]  # transparent green
            world.ccolor = [0, 0, 1, 0.2]  # transparent blue
            selected = world.a
        elif e.key == 'b':
            world.acolor = [1, 0, 0, 0.2]  # transparent red
            world.bcolor = [0, 1, 0, 0.8]  # transparent green
            world.ccolor = [0, 0, 1, 0.2]  # transparent blue
            selected = world.b
        elif e.key == 'c':
            world.acolor = [1, 0, 0, 0.2]  # transparent red
            world.bcolor = [0, 1, 0, 0.4]  # transparent green
            world.ccolor = [0, 0, 1, 0.5]  # transparent blue
            selected = world.c
        elif e.key == 'escape':
            world.acolor = [1, 0, 0, 0.2]  # transparent red
            world.bcolor = [0, 1, 0, 0.4]  # transparent green
            world.ccolor = [0, 0, 1, 0.2]  # transparent blue
            selected = None
        elif e.key == 'up':
            """ move selected 'up' """
            if selected:
                selected.move(Point(0, 0, 0.5))
        elif e.key == 'down':
            """ move selected 'down' """
            if selected:
                selected.move(Point(0, 0, -0.5))
        elif e.key == 'right':
            """ move selected 'right' """
            if selected:
                selected.move(Point(0, 0.5, 0))
        elif e.key == 'left':
            """ move selected 'left' """
            if selected:
                selected.move(Point(0, -0.5, 0))
        else:
            print(e.key, e)
            usage()

    fig.canvas.mpl_connect('key_press_event', on_press)

    def update(angle):
        world.update(3*angle)

        ax.clear()
        ax.set_xlim3d((-10, 10))
        ax.set_ylim3d((-10, 10))
        ax.set_zlim3d((-10, 10))

        world.draw(ax)

    ani = FuncAnimation(fig, update, 1000, interval=100)
    plt.show()

if __name__=='__main__':
    main()