AP_Math: geodesic_grid tool: add option --section

libraries/AP_Math/tools/geodesic_grid/geodesic_grid.py

@@ -20,6 +20,7 @@ import numpy as np
 import sys
 
 import icosahedron as ico
+import grid
 
 def print_code_gen_notice():
     print("/* This was generated with")
@@ -59,7 +60,7 @@ Plot results when applicable.
 )
 
 parser.add_argument(
-    '-s', '--plot-subtriangles',
+    '-b', '--plot-subtriangles',
     action='store_true',
     help="""
 Plot subtriangles as well. This implies -p.
@@ -83,6 +84,17 @@ Get the icosahedron triangle at INDEX.
 """,
 )
 
+parser.add_argument(
+    '-s', '--section',
+    action='append',
+    type=int,
+    nargs='+',
+    help="""
+Get the grid section SECTION. If --plot is passed, then --plot-subtriangles is
+implied.
+""",
+)
+
 parser.add_argument(
     '-u', '--umbrella',
     action='append',
@@ -152,6 +164,23 @@ if args.triangle:
         if args.plot:
             plot.polygon(ico.triangles[i])
 
+if args.section:
+    sections = []
+    for l in args.section:
+        sections += l
+
+    for s in sections:
+        if 0 > s or s >= 4 * len(ico.triangles):
+            print(
+                'Section must be in the range [0,%d)' % 4 * len(ico.triangles),
+                file=sys.stderr,
+            )
+            sys.exit(1)
+        print(grid.section_triangle(s))
+    if args.plot:
+        args.plot_subtriangles = True
+        plot.sections(sections)
+
 if args.umbrella:
     for pivot in args.umbrella:
         for i, x in enumerate(pivot):

libraries/AP_Math/tools/geodesic_grid/grid.py

@@ -0,0 +1,31 @@
+# Copyright (C) 2016  Intel Corporation. All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This file is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+# See the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program.  If not, see <http://www.gnu.org/licenses/>.
+import icosahedron as ico
+
+def section_triangle(s):
+    a, b, c = ico.triangles[s / 4]
+    # project the middle points to the sphere
+    alpha = a.length() / (2.0 * ico.g)
+    ma, mb, mc = alpha * (a + b), alpha * (b + c), alpha * (c + a)
+
+    sub = s % 4
+    if sub == 0:
+        return ico.Triangle(ma, mb, mc)
+    elif sub == 1:
+        return ico.Triangle(a, ma, mc)
+    elif sub == 2:
+        return ico.Triangle(ma, b, mb)
+    else:
+        return ico.Triangle(mc, mb, c)

libraries/AP_Math/tools/geodesic_grid/plot.py

@@ -19,6 +19,7 @@ from mpl_toolkits.mplot3d import Axes3D
 from mpl_toolkits.mplot3d.art3d import Poly3DCollection
 
 import icosahedron as ico
+import grid
 
 fig = plt.figure()
 ax = fig.add_subplot(111, projection='3d')
@@ -36,19 +37,36 @@ ax.invert_xaxis()
 ax.set_aspect('equal')
 
 added_polygons = set()
+added_sections = set()
 
 def polygons(polygons):
     for p in polygons:
         polygon(p)
 
-
 def polygon(polygon):
     added_polygons.add(polygon)
 
+def section(s):
+    added_sections.add(s)
+
+def sections(sections):
+    for s in sections:
+        section(s)
+
 def show(subtriangles=False):
     polygons = []
     facecolors = []
 
+    subtriangle_facecolors = (
+        '#CCCCCC',
+        '#CCE5FF',
+        '#E5FFCC',
+        '#FFCCCC',
+    )
+
+    if added_sections:
+        subtriangles = True
+
     for p in added_polygons:
         try:
             i = ico.triangles.index(p)
@@ -57,23 +75,7 @@ def show(subtriangles=False):
             continue
 
         if subtriangles:
-            a, b, c = p
-
-            # project the middle points to the sphere
-            alpha = a.length() / (2.0 * ico.g)
-            ma, mb, mc = alpha * (a + b), alpha * (b + c), alpha * (c + a)
-
-            polygons.append(ico.Triangle(ma, mb, mc))
-            facecolors.append('#CCCCCC')
-
-            polygons.append(ico.Triangle( a, ma, mc))
-            facecolors.append('#CCE5FF')
-
-            polygons.append(ico.Triangle(ma,  b, mb))
-            facecolors.append('#E5FFCC')
-
-            polygons.append(ico.Triangle(mc, mb,  c))
-            facecolors.append('#FFCCCC')
+            sections(range(i * 4, i * 4 + 4))
         else:
             polygons.append(p)
             facecolors.append('#DDDDDD')
@@ -85,6 +87,11 @@ def show(subtriangles=False):
             mz += z
         ax.text(mx / 2.6, my /2.6, mz / 2.6, i, color='#444444')
 
+    for s in added_sections:
+        subtriangle_index = s % 4
+        polygons.append(grid.section_triangle(s))
+        facecolors.append(subtriangle_facecolors[subtriangle_index])
+
     ax.add_collection3d(Poly3DCollection(
         polygons,
         facecolors=facecolors,
@@ -93,11 +100,9 @@ def show(subtriangles=False):
 
     if subtriangles:
         ax.legend(
-            handles=(
-                mpatches.Patch(color='#CCCCCC', label='Sub-triangle #0'),
-                mpatches.Patch(color='#CCE5FF', label='Sub-triangle #1'),
-                mpatches.Patch(color='#E5FFCC', label='Sub-triangle #2'),
-                mpatches.Patch(color='#FFCCCC', label='Sub-triangle #3'),
+            handles=tuple(
+                mpatches.Patch(color=c, label='Sub-triangle #%d' % i)
+                for i, c in enumerate(subtriangle_facecolors)
             ),
         )