# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

# +
# %matplotlib inline
# %load_ext autoreload
# %autoreload 2
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import scipy as sp
import scipy.stats
import time
import yaml
try:
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader
    

from pydrake.all import (RigidBodyTree, RigidBody)

# +
# These scenes include uniformly randomly distributed nonpenetrating
# object arrangements
DATA_FILE = "data/20180702_uniform_feasible_2_class_10000.yaml"
# These loaded scenes include squares on the left half of the scene, and
# long rectangles on the right half of the scene. All objects
# tend to lie along cardinal directions (theta = 0, pi/2, pi, ...)
#DATA_FILE = "data/20180709_sorted_grid_10000.yaml"

with open(DATA_FILE, "r") as f:
    environments = yaml.load(f, Loader=Loader)
    N_ENVIRONMENTS = len(environments.keys())
    print("Loaded %d environments from file %s" % (N_ENVIRONMENTS, DATA_FILE))

# +
# Make a listing of the observed classes
class_name_to_index_map = {}
class_index_to_name_map = []
n_objects_by_class = {}
counts_by_environment = np.zeros(len(environments.keys()))
current_ind = 0
for i, env_name in enumerate(environments.keys()):
    env = environments[env_name]
    counts_by_environment[i] = env["n_objects"]
    for k in range(env["n_objects"]):
        class_name = env["obj_%04d" % k]["class"]
        if class_name not in class_name_to_index_map:
            class_name_to_index_map[class_name] = current_ind
            class_index_to_name_map.append(class_name)
            current_ind += 1
        if class_name not in n_objects_by_class.keys():
            n_objects_by_class[class_name] = 1
        else:
            n_objects_by_class[class_name] += 1
print class_name_to_index_map, class_index_to_name_map
print n_objects_by_class
N_CLASSES = current_ind

n_hist, n_hist_bins = np.histogram(counts_by_environment,
    bins=range(int(np.ceil(np.max(counts_by_environment)))+2))
n_pdf = n_hist.astype(np.float64)/np.sum(n_hist)
plt.subplot(2, 1, 1)
plt.bar(n_hist_bins[:-1], n_pdf, align="edge")
plt.xlabel("# objects")
plt.ylabel("occurance rate")

plt.subplot(2, 1, 2)
class_pdf = [float(n_objects_by_class[ci]) for ci in class_index_to_name_map]
class_pdf = class_pdf / np.sum(class_pdf)
plt.bar(range(N_CLASSES), class_pdf, align="edge")
plt.xlabel("object class")
plt.ylabel("occurance rate")
plt.tight_layout()

# +
# Draw a few example scenes from the set
import generate_planar_scene_arrangements as psa_utils

def draw_environment(environment, ax):
    rbt, q = psa_utils.build_rbt_from_summary(environment)
    psa_utils.draw_board_state(ax, rbt, q)
    
    patch = patches.Rectangle([0., 0.], 1., 1., fill=True, color=[0., 1., 0.], 
                              linestyle='solid', linewidth=2, alpha=0.3)
    ax.add_patch(patch)

plt.figure().set_size_inches(12, 12)
print "Selection of environments from original distribution"
N = 5
for i in range(N):
    for j in range(N):
        plt.subplot(N, N, i*N+j+1)
        draw_environment(environments["env_%04d" % (i*N+j)], plt.gca())
        plt.grid(True)
plt.tight_layout()
# -

# ## Fitting distributions to generated scene data
#
# As seen in the 20180626 notebook, treating objects as occuring independently loses all notion of object interaction. (Makes sense...) We need to have terms representing placement of pairs of objects.
#
# A distribution we could look at would be the distribution of each object class in the frame of reference of every other object class.

# +
# Preallocate a dictionary of dictionaries
# (i.e. can index [class_name_1][class_name_2] for
# the relative incidences of class_name_1, class_name_2)
relative_positions_by_class_pair = {}
positions_by_class = {}
counts_by_class_pair = {}
counts_by_class = {}
for class_name_1 in class_index_to_name_map:
    relative_positions_by_class = {}
    positions_by_class[class_name_1] = np.zeros((3, n_objects_by_class[class_name_1]))
    counts_by_class[class_name_1] = 0
    for class_name_2 in class_index_to_name_map:
        # TODO(gizatt) I'm overallocating this array by a massive
        # margin...
        relative_positions_by_class_pair[class_name_1, class_name_2] = \
            np.zeros((3, n_objects_by_class[class_name_1]*(n_objects_by_class[class_name_2])))
        counts_by_class_pair[class_name_1, class_name_2] = 0

# Go collect relative poses
for env_name in environments.keys():
    env = environments[env_name]
    # For every object...
    for i in range(env["n_objects"]):
        for j in range(env["n_objects"]):
            o1 = env["obj_%04d" % i]
            ci = o1["class"]
            pose_i = o1["pose"]
            if i == j:
                k = counts_by_class[ci]
                counts_by_class[ci] += 1
                positions_by_class[ci][:, k] = pose_i[:]
            else:
                o2 = env["obj_%04d" % j]
                cj = o2["class"]
                pose_j = o2["pose"]
                # Transform relative pose into pose_i frame
                relative_pose = np.array(pose_j) - np.array(pose_i)
                c, s = np.cos(-pose_i[2]), np.sin(-pose_i[2])
                relative_pose[0:2] = np.array(((c,-s), (s, c))).dot(relative_pose[0:2])
                k = counts_by_class_pair[ci,cj]
                counts_by_class_pair[ci,cj] += 1
                relative_positions_by_class_pair[ci, cj][:, k] = relative_pose
            
for ci in class_index_to_name_map:
    for cj in class_index_to_name_map:
        relative_positions_by_class_pair[ci, cj] = \
            relative_positions_by_class_pair[ci, cj][:, :counts_by_class_pair[ci, cj]]
# -

# Contrary to the 20180702 notebook, I'm going to try GMM instead of KDE. This'll sacrifice quality of representation for size -- with the additional plus that GMMs are composable, to a limited degree.

# +
from composable_gmm import ComposableGMM

class_gmm_fits_by_pair = {}
class_gmm_fits = {}
for i, ci in enumerate(class_index_to_name_map):
    print "Computing GMM for %s" % (ci)
    poses = positions_by_class[ci]
    gmm_fit = ComposableGMM(n_components=20, n_init=3)
    gmm_fit.fit(poses.T)
    class_gmm_fits[ci] = gmm_fit
    for j, cj in enumerate(class_index_to_name_map):
        print "Computing GMM for %s v %s" % (ci, cj)
        relative_poses = relative_positions_by_class_pair[ci, cj]
        gmm_fit = ComposableGMM(n_components=20, n_init=3)
        gmm_fit.fit(relative_poses.T)
        class_gmm_fits_by_pair[ci, cj] = gmm_fit

# -

tslice = 3.1415
plt.figure().set_size_inches(12, 12)
for i, ci in enumerate(class_index_to_name_map):
    ax = plt.subplot(1, N_CLASSES, i + 1)

    xmin = -0.25
    xmax = 1.25
    ymin = -0.25
    ymax = 1.25
    N = 50j
    X, Y = np.mgrid[xmin:xmax:N, ymin:ymax:N]
    positions = np.vstack([X.ravel(), Y.ravel(), np.zeros(X.ravel().shape) + tslice])
    Z = np.reshape(np.exp(class_gmm_fits[ci].score_samples(positions.T)), X.shape)
    h = plt.gca().imshow(np.rot90(Z[:, :]),
                     cmap=plt.cm.gist_earth_r,
                     extent=[xmin, xmax, ymin, ymax],
                     alpha=0.5)
    class_gmm_fits[ci].plot(plt.gca(), facecolor="none")
    poses = positions_by_class[ci]
    theta_min = -0.05
    theta_max = 0.05
    select = np.where(np.logical_and(
        poses[2, :] >= tslice - 0.05,
        poses[2, :] <= tslice + 0.05))
    ax.scatter(poses[0, select], poses[1, select], s=5)
    plt.colorbar(h,fraction=0.046, pad=0.04)
    plt.title("p(%s)" % (ci))
    plt.xlim([xmin, xmax])
    plt.ylim([ymin, ymax])
plt.tight_layout()

tslice = 3.1415
plt.figure().set_size_inches(12, 12)
for i, ci in enumerate(class_index_to_name_map):
    for j, cj in enumerate(class_index_to_name_map):
        ax = plt.subplot(N_CLASSES, N_CLASSES, i*N_CLASSES + j + 1)

        rbt, q = psa_utils.build_rbt_from_summary(
            {"n_objects": 1,
             "obj_0000": {"pose": [0, 0, 0], "class": ci}})
        #psa_utils.draw_board_state(ax, rbt, q)

        xmin = -1.
        xmax = 1.
        ymin = -1.
        ymax = 1.
        N = 50j
        X, Y = np.mgrid[xmin:xmax:N, ymin:ymax:N]
        positions = np.vstack([X.ravel(), Y.ravel(), np.zeros(X.ravel().shape) + tslice])
        Z = np.reshape(np.exp(
                class_gmm_fits_by_pair[ci, cj].score_samples(positions.T)), X.shape)
        h = plt.gca().imshow(np.rot90(Z[:, :]),
                         cmap=plt.cm.gist_earth_r,
                         extent=[xmin, xmax, ymin, ymax],
                         alpha=0.6)
        class_gmm_fits_by_pair[ci, cj].plot(plt.gca(), facecolor="none")

        relative_poses = relative_positions_by_class_pair[ci, cj]
        theta_min = -0.05
        theta_max = 0.05
        select = np.where(np.logical_and(
            relative_poses[2, :] >= tslice - 0.05,
            relative_poses[2, :] <= tslice + 0.05))
        #ax.scatter(relative_poses[0, select], relative_poses[1, select], s=5)
        plt.colorbar(h,fraction=0.046, pad=0.04)
        plt.title("p(%s | %s)" % (cj, ci))
        plt.xlim([xmin, xmax])
        plt.ylim([ymin, ymax])
plt.tight_layout()

# ## Generating arrangements
#
# We'll create new arrangements by:
#
# ~~~python
# initialize_context-info()
# # Sample # of objects from observed object counts
# for i in range(sample_num_objects()):
#     # Sample object class from independent prior over object counts
#     ci = sample_object_class()
#     # Sample a new object given the current context
#     sample_object_position(ci, context)
# ~~~
#
# Sampling the number of objects and determining their classes is done independently across objects (i.e. no dependence of object count or class on each other). We'll want to fix this eventually, but it'll probably be passable.
#
# Sampling object position depends on all previous object positions. This makes an assumption of something like causality -- object placement depends on objects placed before it, but not after it. (Sort of an HMM, I guess? We'll definitely need to work hard (see forward-backward algorithm / inside-out algorithm) for doing inference / evaluating final-state likelihood down the road, if we didn't get to observe the scene being constructed.) Makes generation a lot easier! Specifically, given an arrangement of objects, the object placement probability for the new one is
#
# $$
# p(o_{i+1} | o_{i}, ..., o_{1}, c_{i+1} ..., c_{1}) = {p(o_{i+1}) \large \Pi_{k \in [0, n]}} p(o_{i+1} | c_{i+i}, c_{k}, o_{k})
# $$
# for $k^{th}$ object position $o_k$ and class $c_k$.
#
# We have KDE estimates of $p(o_{i+1})$ or $p(o_{i+1} | c_{i+1}, c_{k}, o_{k})$. To assemble the complete joint distribution, we'll try forming a grid over the space and calculating this value at every point over the grid. Then we can sample from this grid (by normalizing it, picking an arbitrary vectorization, computing its CDF, picking a uniform random value on $[0, 1]$, and taking the first location with CDF $>$ that value.

n_cdf = np.cumsum(n_pdf)
class_cdf = np.cumsum(class_pdf)
do_plotting = True
def sample_object_arrangement(N = 50, do_plotting = False):
    n_objects = np.argmax(n_cdf >= np.random.random())
    env = {}

    tslice = 3.1415
    xmin = -0.25
    xmax = 1.25
    ymin = -0.25
    ymax = 1.25
    X, Y = np.mgrid[xmin:xmax:N, ymin:ymax:N]
    positions = np.vstack([X.ravel(), Y.ravel(), np.zeros(X.ravel().shape) + tslice])


    if do_plotting:
        plt.figure().set_size_inches(16, 12)
    for i in range(n_objects):
        ci = class_index_to_name_map[np.argmax(class_cdf >= np.random.random())]
        total_gmm = class_gmm_fits[ci]

        for k in range(i):
            ok = env["obj_%04d" % k]["pose"]
            ck = env["obj_%04d" % k]["class"]
            c, s = np.cos(-ok[2]), np.sin(-ok[2])
            R = np.array(((c, -s, 0), (s, c, 0), (0, 0, 1)))
            T = np.array(ok)
            total_gmm = total_gmm.prod(total_gmm.transform(T, R), max_components=50)

        # Sample from grid
        new_pose = total_gmm.sample()[0][0, :]
        
        env["n_objects"] = i+1
        env["obj_%04d" % i] = {
            "pose": new_pose.tolist(),
            "class": ci
        }
        if do_plotting:
            print("Object %d: %d gaussians in mixture" % (i, total_gmm.weights_.shape[0]))
            plt.subplot(n_objects, 1, i+1)
            rbt, q = psa_utils.build_rbt_from_summary(env)
            psa_utils.draw_board_state(plt.gca(), rbt, q)
            Z = np.reshape(np.exp(total_gmm.score_samples(positions.T)), X.shape)
            h = plt.gca().imshow(np.rot90(Z[:, :]),
                             cmap=plt.cm.gist_earth_r,
                             extent=[xmin, xmax, ymin, ymax],
                             alpha=0.5)
            
            total_gmm.plot(plt.gca(), facecolor="none")
            plt.xlim(xmin, xmax)
            plt.ylim(ymin, ymax)
            plt.title("Object %d: %s" % (i, ci))
            plt.colorbar(h)

    return env
np.random.seed(1237)
sample_object_arrangement(N = 50, do_plotting = True);
plt.tight_layout()

plt.figure().set_size_inches(16, 16)
for i in range(4):
    for j in range(4):
        env = sample_object_arrangement(N = 10, do_plotting = False)
        print "Generated arrangement (%d, %d) with %d objects" % (i, j, env["n_objects"])
        plt.subplot(4, 4, i*4 + j + 1)
        rbt, q = psa_utils.build_rbt_from_summary(env)
        psa_utils.draw_board_state(plt.gca(), rbt, q)
        patch = patches.Rectangle([0., 0.], 1., 1., fill=True, color=[0., 1., 0.], 
                                  linestyle='solid', linewidth=2, alpha=0.3)
        plt.gca().add_patch(patch)
        plt.xlim(-0.25, 1.25)
        plt.ylim(-0.25, 1.25)
plt.tight_layout()

StarcoderdataJupyterScriptsDedupFiltered