ConstrainedPlanningTorus.py

#!/usr/bin/env python3
 
 
 
# Author: Mark Moll
 
from __future__ import print_function
 
from os.path import dirname, join
import argparse
import math
from functools import partial
import numpy as np
from ConstrainedPlanningCommon import *
 
PI2 = 2 * math.pi
 
# Torus manifold.
 
 
class TorusConstraint(ob.Constraint):
 
    def __init__(self, outer, inner, maze):
        super(TorusConstraint, self).__init__(3, 1)
        self.outer = outer
        self.inner = inner
        self.ppm = ou.PPM()
        self.ppm.loadFile(maze)
 
    def getStartAndGoalStates(self):
        h = self.ppm.getHeight()
        w = self.ppm.getWidth()
 
        for x in range(w):
            for y in range(h):
                p = np.array([x / (w - 1.), y / (h - 1.)], dtype=np.float64)
                c = self.ppm.getPixel(x, y)
                if c.red == 255 and c.blue == 0 and c.green == 0:
                    start = self.mazeToAmbient(p)
                elif c.green == 255 and c.blue == 0 and c.red == 0:
                    goal = self.mazeToAmbient(p)
        return start, goal
 
    def function(self, x, out):
        c = np.array([x[0], x[1], 0])
        nrm = math.sqrt(x[0] * x[0] + x[1] * x[1])
        if not np.isfinite(nrm) or nrm == 0:
            nrm = 1
        out[0] = np.linalg.norm(x - self.outer * c / nrm) - self.inner
 
    def jacobian(self, x, out):
        xySquaredNorm = x[0] * x[0] + x[1] * x[1]
        xyNorm = math.sqrt(xySquaredNorm)
        denom = math.sqrt(x[2] * x[2] + (xyNorm - self.outer)
                          * (xyNorm - self.outer))
        c = (xyNorm - self.outer) * (xyNorm * xySquaredNorm) / \
            (xySquaredNorm * xySquaredNorm * denom)
        out[0, :] = [x[0] * c, x[1] * c, x[2] / denom]
 
    def ambientToMaze(self, x):
        nrm = math.sqrt(x[0] * x[0] + x[1] * x[1])
        h = self.ppm.getHeight()
        w = self.ppm.getWidth()
 
        mx = math.atan2(x[2], nrm - self.outer) / PI2
        mx += (mx < 0)
        mx *= h
        my = math.atan2(x[1], x[0]) / PI2
        my += (my < 0)
        my *= w
        return mx, my
 
    def mazeToAmbient(self, x):
        a = x * PI2
        b = [math.cos(a[0]), 0, math.sin(a[0])] * self.inner
        b[0] += self.outer
 
        norm = math.sqrt(b[0] * b[0] + b[1] * b[1])
        out = np.array([math.cos(a[1]), math.sin(a[1]), 0], dtype=np.float64)
        out *= norm
        out[2] = b[2]
        return out
 
    def mazePixel(self, x):
        h = self.ppm.getHeight()
        w = self.ppm.getWidth()
 
        if x[0] < 0 or x[0] >= w or x[1] < 0 or x[1] >= h:
            return False
 
        c = self.ppm.getPixel(int(x[0]), int(x[1]))
        return not (c.red == 0 and c.blue == 0 and c.green == 0)
 
    def isValid(self, state):
        return self.mazePixel(self.ambientToMaze(state))
 
 
def torusPlanningBench(cp, planners):
    print(planners)
    cp.setupBenchmark(planners, "torus")
    cp.runBenchmark()
 
 
def torusPlanningOnce(cp, planner, output):
    cp.setPlanner(planner)
 
    # Solve the problem
    stat = cp.solveOnce(output, "torus")
 
    if output:
        ou.OMPL_INFORM("Dumping problem information to `torus_info.txt`.")
        with open("torus_info.txt", "w") as infofile:
            print(cp.spaceType, file=infofile)
 
    cp.atlasStats()
 
    if output:
        cp.dumpGraph("torus")
 
    return stat
 
 
def torusPlanning(options):
    # Create the ambient space state space for the problem.
    rvss = ob.RealVectorStateSpace(3)
 
    bounds = ob.RealVectorBounds(3)
    bounds.setLow(-(options.outer + options.inner))
    bounds.setHigh(options.outer + options.inner)
 
    rvss.setBounds(bounds)
 
    # Create our constraint.
    constraint = TorusConstraint(options.outer, options.inner, options.maze)
 
    cp = ConstrainedProblem(options.space, rvss, constraint, options)
 
    start, goal = constraint.getStartAndGoalStates()
    print("Start = ", start)
    print("Goal = ", goal)
 
    sstart = ob.State(cp.css)
    sgoal = ob.State(cp.css)
    for i in range(3):
        sstart[i] = start[i]
        sgoal[i] = goal[i]
    cp.setStartAndGoalStates(sstart, sgoal)
    cp.ss.setStateValidityChecker(ob.StateValidityCheckerFn(partial(
        TorusConstraint.isValid, constraint)))
 
    planners = options.planner.split(",")
    if not options.bench:
        torusPlanningOnce(cp, planners[0], options.output)
    else:
        torusPlanningBench(cp, planners)
 
if __name__ == "__main__":
    defaultMaze = join(join(dirname(__file__), "mazes"), "normal.ppm")
    parser = argparse.ArgumentParser()
    parser.add_argument("-o", "--output", action="store_true",
                        help="Dump found solution path (if one exists) in plain text and planning "
                        "graph in GraphML to `torus_path.txt` and `torus_graph.graphml` "
                        "respectively.")
    parser.add_argument("--bench", action="store_true",
                        help="Do benchmarking on provided planner list.")
    parser.add_argument("--outer", type=float, default=2,
                        help="Outer radius of torus.")
    parser.add_argument("--inner", type=float, default=1,
                        help="Inner radius of torus.")
    parser.add_argument("--maze", default=defaultMaze,
                        help="Filename of maze image (in .ppm format) to use as obstacles on the "
                        "surface of the torus.")
    addSpaceOption(parser)
    addPlannerOption(parser)
    addConstrainedOptions(parser)
    addAtlasOptions(parser)
 
    torusPlanning(parser.parse_args())