refactor as a package

sho/plot.py

@@ -1,17 +1,31 @@
 import numpy as np
 from matplotlib import cm
+import matplotlib.pyplot as plt
 import itertools
+from mpl_toolkits.mplot3d import Axes3D
+
+from . import x,y,distance
+
+def sphere(x,offset=0.5):
+    """Computes the square of a multi-dimensional vector x."""
+    f = 0
+    for i in range(len(x)):
+        f += (x[i]-offset)**2
+    return -1 * f
+
 
 def surface(ax, shape, f):
     Z = np.zeros( shape )
     for y in range(shape[0]):
         for x in range(shape[1]):
-            Z[y][x] = f( (x/shape[0],y/shape[1]), 0.5 )
+            Z[y][x] = f( (x,y), shape[0]/2 )
 
     X = np.arange(0,shape[0],1)
     Y = np.arange(0,shape[1],1)
     X,Y = np.meshgrid(X,Y)
-    ax.plot_surface(X, Y, Z, cmap=cm.viridis)
+    #ax.plot_surface(X, Y, Z, cmap=cm.viridis)
+    ax.plot_surface(X, Y, Z)
+
 
 def path(ax, shape, history):
     def pairwise(iterable):
@@ -21,11 +35,11 @@ def path(ax, shape, history):
 
     k=0
     for i,j in pairwise(range(len(history)-1)):
-        xi = history[i][1][0]*shape[0]
-        yi = history[i][1][1]*shape[1]
+        xi = history[i][1][0]
+        yi = history[i][1][1]
         zi = history[i][0]
-        xj = history[j][1][0]*shape[0]
-        yj = history[j][1][1]*shape[1]
+        xj = history[j][1][0]
+        yj = history[j][1][1]
         zj = history[j][0]
         x = [xi, xj]
         y = [yi, yj]
@@ -33,3 +47,52 @@ def path(ax, shape, history):
         ax.plot(x,y,z, color=cm.RdYlBu(k))
         k+=1
 
+
+def highlight_sensors(domain, sensors, val=2):
+    """Add twos to the given domain, in the cells where the given
+    sensors are located.
+
+    >>> highlight_sensors( [[0,0],[1,1]], [(0,0),(1,1)] )
+    [[2, 0], [1, 2]]
+    """
+    for s in sensors:
+        # `coverage` fills the domain with ones,
+        # adding twos will be visible in an image.
+        domain[y(s)][x(s)] = val
+    return domain
+
+
+if __name__=="__main__":
+    import snp
+
+    w = 100
+    shape = (w,w)
+    history = []
+
+    val,sol = snp.greedy(
+            snp.make_func(sphere,
+                offset = w/2),
+            snp.make_init(snp.num_rand,
+                dim = 2 * 1,
+                scale = w),
+            snp.make_neig(snp.num_neighb_square,
+                scale = w/10),
+            snp.make_iter(
+                    snp.several,
+                    agains = [
+                        snp.make_iter(snp.iter_max,
+                            nb_it = 100),
+                        snp.make_iter(snp.history,
+                            history = history)
+                    ]
+                )
+        )
+    sensors = snp.num_to_sensors(sol)
+
+    #print("\n".join([str(i) for i in history]))
+
+    fig = plt.figure()
+    ax = fig.gca(projection='3d')
+    surface(ax, shape, sphere)
+    path(ax, shape, history)
+    plt.show()