initial commit

apps/geometry.py

@@ -0,0 +1,226 @@
+import math
+import torch
+
+class GeometryLoss:
+    def __init__(self, pathObj, xyalign=True, parallel=True, smooth_node=True):
+        self.pathObj=pathObj
+        self.pathId=pathObj.id
+        self.get_segments(pathObj)
+        if xyalign:
+            self.make_hor_ver_constraints(pathObj)
+
+        self.xyalign=xyalign
+        self.parallel=parallel
+        self.smooth_node=smooth_node
+
+        if parallel:
+            self.make_parallel_constraints(pathObj)
+
+        if smooth_node:
+            self.make_smoothness_constraints(pathObj)
+
+    def make_smoothness_constraints(self,pathObj):
+        self.smooth_nodes=[]
+        for idx, node in enumerate(self.iterate_nodes()):
+            sm, t0, t1=self.node_smoothness(node,pathObj)
+            if abs(sm)<1e-2:
+                self.smooth_nodes.append((node,((t0.norm()/self.segment_approx_length(node[0],pathObj)).item(),(t1.norm()/self.segment_approx_length(node[1],pathObj)).item())))
+                #print("Node {} is smooth (smoothness {})".format(idx,sm))
+            else:
+                #print("Node {} is not smooth (smoothness {})".format(idx, sm))
+                pass
+
+    def node_smoothness(self,node,pathObj):
+        t0=self.tangent_out(node[0],pathObj)
+        t1=self.tangent_in(node[1],pathObj)
+        t1rot=torch.stack((-t1[1],t1[0]))
+        smoothness=t0.dot(t1rot)/(t0.norm()*t1.norm())
+
+        return smoothness, t0, t1
+
+    def segment_approx_length(self,segment,pathObj):
+        if segment[0]==0:
+            #line
+            idxs=self.segList[segment[0]][segment[1]]
+            #should have a pair of indices now
+            length=(pathObj.points[idxs[1],:]-pathObj.points[idxs[0],:]).norm()
+            return length
+        elif segment[0]==1:
+            #quadric
+            idxs = self.segList[segment[0]][segment[1]]
+            # should have a pair of indices now
+            length = (pathObj.points[idxs[1],:] - pathObj.points[idxs[0],:]).norm()+(pathObj.points[idxs[2],:] - pathObj.points[idxs[1],:]).norm()
+            return length
+        elif segment[0]==2:
+            #cubic
+            idxs = self.segList[segment[0]][segment[1]]
+            # should have a pair of indices now
+            length = (pathObj.points[idxs[1],:] - pathObj.points[idxs[0],:]).norm()+(pathObj.points[idxs[2],:] - pathObj.points[idxs[1],:]).norm()+(pathObj.points[idxs[3],:] - pathObj.points[idxs[2],:]).norm()
+            return length
+
+    def tangent_in(self, segment,pathObj):
+        if segment[0]==0:
+            #line
+            idxs=self.segList[segment[0]][segment[1]]
+            #should have a pair of indices now
+            tangent=(pathObj.points[idxs[1],:]-pathObj.points[idxs[0],:])/2
+            return tangent
+        elif segment[0]==1:
+            #quadric
+            idxs = self.segList[segment[0]][segment[1]]
+            # should have a pair of indices now
+            tangent = (pathObj.points[idxs[1],:] - pathObj.points[idxs[0],:])
+            return tangent
+        elif segment[0]==2:
+            #cubic
+            idxs = self.segList[segment[0]][segment[1]]
+            # should have a pair of indices now
+            tangent = (pathObj.points[idxs[1],:] - pathObj.points[idxs[0],:])
+            return tangent
+
+        assert(False)
+
+    def tangent_out(self, segment, pathObj):
+        if segment[0] == 0:
+            # line
+            idxs = self.segList[segment[0]][segment[1]]
+            # should have a pair of indices now
+            tangent = (pathObj.points[idxs[0],:] - pathObj.points[idxs[1],:]) / 2
+            return tangent
+        elif segment[0] == 1:
+            # quadric
+            idxs = self.segList[segment[0]][segment[1]]
+            # should have a pair of indices now
+            tangent = (pathObj.points[idxs[1],:] - pathObj.points[idxs[2],:])
+            return tangent
+        elif segment[0] == 2:
+            # cubic
+            idxs = self.segList[segment[0]][segment[1]]
+            # should have a pair of indices now
+            tangent = (pathObj.points[idxs[2],:] - pathObj.points[idxs[3],:])
+            return tangent
+
+        assert (False)
+
+    def get_segments(self, pathObj):
+        self.segments=[]
+        self.lines = []
+        self.quadrics=[]
+        self.cubics=[]
+        self.segList =(self.lines,self.quadrics,self.cubics)
+        idx=0
+        total_points=pathObj.points.shape[0]
+        for ncp in pathObj.num_control_points.numpy():
+            if ncp==0:
+                self.segments.append((0,len(self.lines)))
+                self.lines.append((idx, (idx + 1) % total_points))
+                idx+=1
+            elif ncp==1:
+                self.segments.append((1, len(self.quadrics)))
+                self.quadrics.append((idx, (idx + 1), (idx+2) % total_points))
+                idx+=ncp+1
+            elif ncp==2:
+                self.segments.append((2, len(self.cubics)))
+                self.cubics.append((idx, (idx + 1), (idx+2), (idx + 3) % total_points))
+                idx += ncp + 1
+
+    def iterate_nodes(self):
+        for prev, next in zip([self.segments[-1]]+self.segments[:-1],self.segments):
+            yield (prev, next)
+
+    def make_hor_ver_constraints(self, pathObj):
+        self.horizontals=[]
+        self.verticals=[]
+        for idx, line in enumerate(self.lines):
+            startPt=pathObj.points[line[0],:]
+            endPt=pathObj.points[line[1],:]
+
+            dif=endPt-startPt
+
+            if abs(dif[0])<1e-6:
+                #is horizontal
+                self.horizontals.append(idx)
+
+            if abs(dif[1])<1e-6:
+                #is vertical
+                self.verticals.append(idx)
+
+    def make_parallel_constraints(self,pathObj):
+        slopes=[]
+        for lidx, line in enumerate(self.lines):
+            startPt = pathObj.points[line[0], :]
+            endPt = pathObj.points[line[1], :]
+
+            dif = endPt - startPt
+
+            slope=math.atan2(dif[1],dif[0])
+            if slope<0:
+                slope+=math.pi
+
+            minidx=-1
+            for idx, s in enumerate(slopes):
+                if abs(s[0]-slope)<1e-3:
+                    minidx=idx
+                    break
+
+            if minidx>=0:
+                slopes[minidx][1].append(lidx)
+            else:
+                slopes.append((slope,[lidx]))
+
+        self.parallel_groups=[sgroup[1] for sgroup in slopes if len(sgroup[1])>1 and (not self.xyalign or (sgroup[0]>1e-3 and abs(sgroup[0]-(math.pi/2))>1e-3))]
+
+    def make_line_diff(self,pathObj,lidx):
+        line = self.lines[lidx]
+        startPt = pathObj.points[line[0], :]
+        endPt = pathObj.points[line[1], :]
+
+        dif = endPt - startPt
+        return dif
+
+    def calc_hor_ver_loss(self,loss,pathObj):
+        for lidx in self.horizontals:
+            dif = self.make_line_diff(pathObj,lidx)
+            loss+=dif[0].pow(2)
+
+        for lidx in self.verticals:
+            dif = self.make_line_diff(pathObj,lidx)
+            loss += dif[1].pow(2)
+
+    def calc_parallel_loss(self,loss,pathObj):
+        for group in self.parallel_groups:
+            diffs=[self.make_line_diff(pathObj,lidx) for lidx in group]
+            difmat=torch.stack(diffs,1)
+            lengths=difmat.pow(2).sum(dim=0).sqrt()
+            difmat=difmat/lengths
+            difmat=torch.cat((difmat,torch.zeros(1,difmat.shape[1])))
+            rotmat=difmat[:,list(range(1,difmat.shape[1]))+[0]]
+            cross=difmat.cross(rotmat)
+            ploss=cross.pow(2).sum()*lengths.sum()*10
+            loss+=ploss
+
+    def calc_smoothness_loss(self,loss,pathObj):
+        for node, tlengths in self.smooth_nodes:
+            sl,t0,t1=self.node_smoothness(node,pathObj)
+            #add smoothness loss
+            loss+=sl.pow(2)*t0.norm().sqrt()*t1.norm().sqrt()
+            tl=((t0.norm()/self.segment_approx_length(node[0],pathObj))-tlengths[0]).pow(2)+((t1.norm()/self.segment_approx_length(node[1],pathObj))-tlengths[1]).pow(2)
+            loss+=tl*10
+
+    def compute(self, pathObj):
+        if pathObj.id != self.pathId:
+            raise ValueError("Path ID {} does not match construction-time ID {}".format(pathObj.id,self.pathId))
+
+        loss=torch.tensor(0.)
+        if self.xyalign:
+            self.calc_hor_ver_loss(loss,pathObj)
+
+        if self.parallel:
+            self.calc_parallel_loss(loss, pathObj)
+
+        if self.smooth_node:
+            self.calc_smoothness_loss(loss,pathObj)
+
+        #print(loss.item())
+
+        return loss