# pyCFI Development Notebook - Improved Full 3D Approach

### Notes

- Open issues are marked in comments using the word `FLAG`

### Prep

```python
### Imports

import itertools, collections

import numpy as np
import matplotlib.pyplot as plt
import networkx as nx

from skimage import io
from scipy import spatial
from scipy import interpolate
from scipy import optimize
import sympy as sym

from ipywidgets import interact
from mpl_toolkits.mplot3d import Axes3D
```

```python
### Parameters & Settings

fpath = '../Data/Generated/three_intersecting_spheres_aniso.tif'
res   = np.array([1.0, 0.5, 0.5])  # Voxel sizes (z,y,x) in microns
```

```python
### Load input segmentation stack

im = io.imread(fpath)
print(im.dtype, im.shape)
```

```python
### Show input stack

@interact(z=(0, im.shape[0]-1, 1))
def show_stack(z=im.shape[0]//2):
    plt.figure(figsize=(8,8))
    plt.imshow(im[z], cmap='gray')
    plt.show()
```

### Identifying Object Outlines

```python
### Identify outline voxels by comparing shifted images

# Pad the image by 1 voxel on all sides
im_pad = np.pad(im, 1, mode='reflect')

# Get possible shifts in all directions
shifts = itertools.product([0,1], repeat=3)

# Check and accumulate differences in shifts
outlines = np.zeros_like(im, dtype=np.bool)
for shift in shifts:
    zs0, ys0, xs0 = [slice(1, None) if s else slice(None) for s in shift]
    zs1, ys1, xs1 = [slice(None,-1) if s else slice(None) for s in shift]
    comparison = im_pad[zs0, ys0, xs0] != im_pad[zs1, ys1, xs1]
    outlines  += comparison[:im.shape[0],  :im.shape[1],  :im.shape[2]]
    outlines  += comparison[-im.shape[0]:, -im.shape[1]:, -im.shape[2]:]  # Symmetry

# Re-annotate the cell identities
outlines_id = outlines * im

# Report
print(outlines.dtype, outlines.shape)
print(outlines_id.dtype, outlines_id.shape)
```

```python
### Show identified outlines

@interact(z=(0, im.shape[0]-1, 1))
def show_stack(z=im.shape[0]//2):
    plt.figure(figsize=(8,8))
    plt.imshow(outlines_id[z], cmap='gray')
    plt.show()
```

### Identifying Triple Nodes (TNs) and Triple Junctions (TJs)

```python
### Find coordinates of all voxels involved in triple junctions

# FLAG: PERFORMANCE -- This could potentially be done with image shifting much like `outlines` above!
# FLAG: PRECISION -- For the coordinates, would a `+1.0` be more appropriate at interfaces between two cells?

# Get Outline Indices (OIs) and Outline Coordinates (OCs)
OIs = np.array(np.where(outlines)).T
OCs = (OIs + 0.5) * res

# Go through OIs and find TN Indices (TNIs)
TNIs = []
for OI in OIs:
    selection = im_pad[OI[0]+1:OI[0]+3, OI[1]+1:OI[1]+3, OI[2]+1:OI[2]+3]
    if len(set(selection.flatten())) == 3:
        TNIs.append(OI)
TNIs = np.array(TNIs)

# Convert to TN Coordinates (TNCs)
TNCs = (TNIs + 0.5) * res

# Report
print('OCs: ', OCs.shape)
print('TNIs:', TNIs.shape)
print('TNCs:', TNCs.shape)
```

```python
### Build a dict of TJs structured as: {tuple(cell1_ID, cell2_ID, cell3_ID) : array(INDICES INTO TNIs/TNCs)}

# Prepare defaultdict
TJs = collections.defaultdict(lambda : [])

# Go through TNs, create IDs, assign coordinates to IDs
for idx,TNI in enumerate(TNIs):
    selection = im_pad[np.int(TNI[0])+1:np.int(TNI[0])+3,
                       np.int(TNI[1])+1:np.int(TNI[1])+3,
                       np.int(TNI[2])+1:np.int(TNI[2])+3]
    TJ_ID = tuple(sorted(set(selection.flatten())))
    TJs[TJ_ID].append(idx)

# Convert TJ lists to numpy arrays
for key in TJs.keys():
    TJs[key] = np.array(TJs[key])
```

```python
### Show identified TJs on image stack

@interact(z=(0, im.shape[0]-1, 1))
def show_stack(z=im.shape[0]//2):

    # Prep and plot image
    plt.figure(figsize=(8,8))
    plt.imshow(outlines_id[z], cmap='gray')

    # For each TJ...
    for TJ_num,TJ_ID in enumerate(TJs.keys()):

        # Get the TJ's TNs in the selected z plane
        TNs_in_plane = TNIs[TJs[TJ_ID]][TNIs[TJs[TJ_ID]][:,0]==z]

        # Plot the points
        plt.scatter(TNs_in_plane[:, 2], TNs_in_plane[:, 1],
                    c=[TJ_num for _ in range(TNs_in_plane.shape[0])], # Coloring trick!
                    cmap='hsv', vmin=0, vmax=len(TJs), s=20)

    # Finish
    plt.show()
```

```python
### Show identified TJs as 3D scatter

# Prepare the plot
fig = plt.figure(figsize=(12,12))
ax  = fig.add_subplot(111, projection='3d')

# Plot each TJ in a different color
for TJ_num,TJ_ID in enumerate(TJs.keys()):
    ax.scatter(TNCs[TJs[TJ_ID]][:,2], TNCs[TJs[TJ_ID]][:,1], TNCs[TJs[TJ_ID]][:,0],
               c=[TJ_num for _ in range(TJs[TJ_ID].shape[0])],
               cmap='hsv', vmin=0, vmax=len(TJs), s=10)

## Also show cell outlines [may take several seconds to render!]
#ax.scatter([c[2] for c in OCs],
#           [c[1] for c in OCs],
#           [c[0] for c in OCs],
#           c='gray', alpha=0.01, linewidth=0, s=5)

## Axis limits
#ax.set_xlim([0,200])
#ax.set_ylim([0,200])
#ax.set_zlim([0,200])

# Finish
plt.show()
```

### Identifying Double Nodes (DNs) and Double Junctions (DJs)

```python
### Find coordinates of all voxels involved in DOUBLE junctions

# FLAG: PERFORMANCE -- Same as for TNI/TNC extraction above!
# FLAG: PRECISION -- Same as for TNI/TNC extraction above!

# Go through OIs and find DN Indices (DNIs)
DNIs = []
for OI in OIs:
    selection = im_pad[OI[0]+1:OI[0]+3,
                       OI[1]+1:OI[1]+3,
                       OI[2]+1:OI[2]+3]
    if len(set(selection.flatten())) == 2:
        DNIs.append(OI)
DNIs = np.array(DNIs)

# Convert to DN Coordinates (DNCs)
DNCs = (DNIs + 0.5) * res

# Report
print('OCs: ', OCs.shape)
print('DNIs:', DNIs.shape)
print('DNCs:', DNCs.shape)
```

```python
### Build a dict of Double Junctions (DJs) structured as: {tuple(cell1_ID, cell2_ID) : array(INDICES INTO DNIs/DNCs)}

# Prepare defaultdict
DJs = collections.defaultdict(lambda : [])

# Go through DNs, create IDs, assign coordinates to IDs
for idx,DNI in enumerate(DNIs):
    selection = im_pad[np.int(DNI[0])+1:np.int(DNI[0])+3,
                       np.int(DNI[1])+1:np.int(DNI[1])+3,
                       np.int(DNI[2])+1:np.int(DNI[2])+3]
    DJ_ID = tuple(sorted(set(selection.flatten())))
    DJs[DJ_ID].append(idx)

# Convert DJ lists to numpy arrays
for key in DJs.keys():
    DJs[key] = np.array(DJs[key])
```

```python
### Show identified DJs on image stack

@interact(z=(0, im.shape[0]-1, 1))
def show_stack(z=im.shape[0]//2):

    # Prep and plot image
    plt.figure(figsize=(8,8))
    plt.imshow(outlines_id[z], cmap='gray')

    # For each DJ...
    for DJ_num,DJ_ID in enumerate(DJs.keys()):

        # Get the DJ's DNs in the selected z plane
        DNs_in_plane = DNIs[DJs[DJ_ID]][DNIs[DJs[DJ_ID]][:,0]==z]

        # Plot the points
        plt.scatter(DNs_in_plane[:, 2], DNs_in_plane[:, 1],
                    c=[DJ_num for _ in range(DNs_in_plane.shape[0])], # Coloring trick!
                    cmap='hsv', vmin=0, vmax=len(DJs), s=5, lw=0, alpha=0.5)

    # Finish
    plt.show()
```

### Fitting Splines to TJs

Turns out spline fitting requires the input points to be roughly in order along the spline, which isn't guaranteed in our case. Ordering the points happens to be far harder problem than one might imagine (it's a variation of traveling salesman) but luckily it can be solved quite well with a Breadth-First Search (BFS). This solution is partially inspired by Imanol Luengo's answer to [this SO question](https://stackoverflow.com/questions/37742358/sorting-points-to-form-a-continuous-line).

**Warning 1:** This will fail for geometries that exhibit "crossings" or "forks" of any kind. Although that should be very rare/non-existent in the data, a special form of "fork" is the circle. In case of a fully circular TJ, which occurs when two cells neatly touch each other, this will fail (unless some points are removed from the TJ). I couldn't come up with a way of fixing this but divised the `InvalidPathError` to at least pick up on such cases. However, **it may be too stringent** as it is currently implemented!

**Warning 2:** Simply rescaling the z axis a little bit already led to renewed problems with this approach, so I'm starting to seriously doubt its robustness. We'll have to keep a close eye on this and possibly somehow develop a better solution if problems keep cropping up. Maybe some sort of modified graph search (rather than straight up BFS) would be a possibility...

```python
### Function to reorder TEs along the progression of the TJ

# FLAG: ROBUSTNESS -- I still have my doubts as to the robustness of this approach (see warnings above)!
#                     I keep wondering if there isn't a better way!

def sort_line_coords(coords, N_neighbors=10, source=None,
                     return_argsort=False, ignore_path_check=False):
    """Given a set of coordinates that roughly lie on a 1D curve in mD space
    (but may be in random order), sort the points such that they roughly follow
    the curve's progression.

    Uses a breadth-first search tree on a nearest-neighbor graph of the coords,
    which happens to result in the best possible sort. Does not work as intended
    for closed curves and curves that form any kind of fork or crossing; an
    Error is raised in such cases.

    Parameters
    ----------
    coords : array of shape (N_points, M_dimensions)
        Coordinates of points roughly lying on a point in M-dimensional space.
    N_neighbors : int, optional, default 10
        Number of nearest neighbors to include for each graph. If this is set
        too low, connected components may form and no complete solution is
        possible (raises an Exception). If this is set too high, the resulting
        sort is very imprecises. The ideal value must be determined empirically.
        When used to prepare TJs for spline fitting in the context of pyCFI, the
        default (10) is a reasonably choice and the outcome is largely robust
        to changes between values of 5 and 20.
    source : None or int, optional, default None
        The source is a point at one of the two ends of the line. If None, the
        point is automatically determined by testing all different points and
        selecting the one that yields the best sort (by minimizing the resulting
        path distance). If source is an int, it indexes into coords to specify
        the end point from which the sort is constructed. This saves a lot of
        time compared to the automated search, especially if there are many
        points, however it requires prior knowledge of the end point.
    return_argsort : bool, optional, default False
        If True, the index array that sorts the points into the best order is
        returned as a second result. Otherwise, only a sorted version of coords
        is returned.
    ignore_path_check : bool, optional, default False
        If True, the final path is not cross-checked and no InvalidPathErrors
        can be raised (see Exceptions below).

    Returns
    -------
    sorted_coords : array of shape (N_points, M_dimensions)
        The same set of points as in the input coords but sorted along the
        curve's progression in space.
    best_path : array of shape (N_points,)

Index array that sorts points along the curve's progression in space.
        Only returned if return_argsort is set to True.

    Exceptions
    ----------
    InvalidPathError : If the curve is closed or contains forks/crossings, the
        sort fails, which is reflected in the fact that the final path will
        contain steps that do not have corresponding edges on the graph. In
        this case, InvalidPathError is raised. This may also occur under other
        dubious circumstances, e.g. if the input data is not a curve at all
        or if it is a very broad curve or if N_neighbors is too low.
    """

    # Get pairwise distances (if needed)
    if source is None:
        dists = spatial.distance.squareform(spatial.distance.pdist(coords))

    # Get nearest neighbors
    kdtree  = spatial.cKDTree(coords)
    _, KNNs = kdtree.query(coords, k=N_neighbors)

    # Build adjacency matrix
    adj_M = np.zeros((coords.shape[0], coords.shape[0]), dtype=np.bool)
    for i,N in enumerate(KNNs):
        adj_M[i,N] = True

    # Construct networkx graph
    G = nx.from_numpy_array(adj_M)
    if not nx.is_connected(G):
        raise Exception('sort_line_coords: adjacency graph is not fully connected!')

    # If a source node is given, just get its BFS tree
    if source is not None:
        best_path = list(nx.bfs_tree(G, source))

    # Otherwise, find the best BFS tree from all sources
    if source is None:
        paths = []
        costs = []
        for n in G.nodes():

            # Get BFS tree
            path = list(nx.bfs_tree(G, n))

            # Get sum of all distances within tree
            cost = 0.0
            for n0,n1 in zip(path, path[1:]):
                cost += dists[n0, n1]

            # Keep results
            paths.append(path)
            costs.append(cost)

        # Select the best solution
        best_path = paths[np.argmin(costs)]

    # Test for cases that probably failed
    if not ignore_path_check:
        for p1,p2 in zip(best_path, best_path[1:]):
            if not G.has_edge(p1,p2):
                class InvalidPathError(Exception): pass
                raise InvalidPathError("The sort path uses an edge that is not on the graph. "+
                                       "This should not happen and probably implies that the "+
                                       "curve is cyclical or has a fork/crossing.")

    # Sort coords and return
    if return_argsort:
        return coords[best_path], best_path
    else:
        return coords[best_path]
```

```python
### A quick test of the TJ sorting

# FLAG: ROBUSTNESS -- Currently, ingore_path_check has to be set to True for this to work
#                     when z is properly rescaled (although the sort overall actually
#                     doesn't look too bad).

# Grab a TJ and compute the sort
TJCs = TNCs[TJs[list(TJs.keys())[0]]]
sorted_TJCs = sort_line_coords(TJCs, ignore_path_check=True)

# Get sorted pairwise distances
sorted_dists = spatial.distance.squareform(spatial.distance.pdist(sorted_TJCs))

# Plot pairwise distances
fig, ax = plt.subplots(1, 2, figsize=(13,6.5))
ax[0].imshow(sorted_dists)
ax[0].set_title("sorted distances")
ax[0].set_xlabel('TNs'); ax[0].set_ylabel('TNs')

# Plot path/tree on scatter
ax[1].scatter(TJCs[:,1], TJCs[:,0], s=20, alpha=0.5)
ax[1].plot(sorted_TJCs[:,1], sorted_TJCs[:,0], c='r', alpha=0.75)
ax[1].set_title("path scatter")
ax[1].set_xlabel('y'); ax[1].set_ylabel('z')

# Done
plt.tight_layout()
plt.show()
```

```python
### Wrapper for spline fitting

def wrap_splprep(coords, k=3, verbose=False):
    """Fit an nD spline with scipy.interpolate.splprep.

    coords : array (points, dimensions) : input data
    k=3 : integer : degrees of freedom
    verbose=False : bool : wether to print all outputs

    returns -> tck : tuple (knots, coefficients, k) :
               fit parameters as used by splev
    """

    # Fit the spline and unpack the (weirdly packaged) results
    tcku, fp, ier, msg = interpolate.splprep(coords.T, k=k, full_output=True)
    tck, u = tcku

    # Report the results
    if verbose:
        print ('\nt (knots, tck[0]):\n' , tck[0])
        print ('\nc (coefficients, tck[1]):\n' , tck[1])
        print ('\nk (degree, tck[2]):' , tck[2])
        print ('\nu (evaluation points):\n', u)
        print ('\nfp (residual error):', fp)
        print ('\nier (error code; success is ier<=0):', ier)
        print ('\nmsg (message from FITPACK):\n', msg)

    # Raise an error if FITPACK indicates failure
    if ier > 0:
        raise Exception('ier is >0, indicating that FITPACK failed somehow. '+
                        'The message from FITPACK was:\n'+msg)

    # Return the only result relevant to spline evaluation
    return tck
```

```python
### Perform sorting and spline fitting on all TJs

# Parameters
num_ts = 100    # Determines the number of TNs that will be analyzed throughout the rest of the pipeline
tng_dv = 10e-2  # FLAG -- PRECISION: Should this be smaller? FLAG -- ROBUSTNESS: Should this scale with res?

# Output dicts
TJs_spline_tck     = {}  # Fitted splines for each TJ
TJs_spline_t       = {}  # Parameter (t) values for evaluation
TJs_spline_ev      = {}  # Evaluated splines (at each t) for each TJ
TJs_spline_tangent = {}  # Tangents to splines for each TJ

# For each TJ...
for TJ_ID in TJs.keys():

    # Sort coordinates along the line
    sorted_TJCs, TJ_argsort = sort_line_coords(TNCs[TJs[TJ_ID]],
                                               return_argsort=True,
                                               ignore_path_check=True)
    TJs[TJ_ID] = TJs[TJ_ID][TJ_argsort]

    # Perform spline fitting
    tck = wrap_splprep(sorted_TJCs)
    TJs_spline_tck[TJ_ID] = tck

    # Evaluate the spline in 1000 regular intervals
    TJs_spline_t[TJ_ID] = np.linspace(0.0, 1.0, num_ts)
    ev = interpolate.splev(TJs_spline_t[TJ_ID], tck)
    ev = np.array(ev).T
    TJs_spline_ev[TJ_ID] = ev

    # Also evaluate with slight deviation forward and backward
    evD1 = np.array(interpolate.splev(TJs_spline_t[TJ_ID]+tng_dv, tck)).T
    evD2 = np.array(interpolate.splev(TJs_spline_t[TJ_ID]-tng_dv, tck)).T

    # Approximate the tangent vector as the sum of the deviatory vectors
    tangent_vec = ((evD1 - ev) + (ev - evD2)) / 2.0
    TJs_spline_tangent[TJ_ID] = tangent_vec
```

```python
### Visualize the fitted splines and the tangent vectors as 3D scatter

# Prepare the plot
fig = plt.figure(figsize=(12,12))
ax  = fig.add_subplot(111, projection='3d')

# Plot each TJ spline
for TJ_num,TJ_ID in enumerate(TJs.keys()):
    ax.plot(TJs_spline_ev[TJ_ID][:,2],
            TJs_spline_ev[TJ_ID][:,1],
            TJs_spline_ev[TJ_ID][:,0],
            lw=3)

# Add the tangent vectors
for TJ_num, TJ_ID in enumerate(TJs.keys()):
    for splpt, tngvec in zip(TJs_spline_ev[TJ_ID][::5], TJs_spline_tangent[TJ_ID][::5]):
        plt.plot([splpt[2], splpt[2]+tngvec[2]],
                 [splpt[1], splpt[1]+tngvec[1]],
                 [splpt[0], splpt[0]+tngvec[0]],
                 'r-', alpha=0.5)

## Axis limits
#ax.set_xlim([0,200])
#ax.set_ylim([0,200])
#ax.set_zlim([0,200])

# Finish
plt.show()
```

### Projecting DNs onto TJ-Orthogonal Dihedral Planes

```python
### Sympy function to project close-by outline points onto a TN's TJ-orthogonal plane

# FLAG: PERFORMANCE -- Save the resulting numpy func so that the symbolic solving doesn't need
#                      to be rerun each time the code is executed! This is probably best done
#                      by copying the function out into a .py file and importing it from there.
#                      In the process, axis keywords could perhaps be added to handle vectorized
#                      execution across many points/planes (see issue flag below).
# FLAG: ROBUSTNESS -- Simply doing Gram-Schmidt as we currently do does not preserve the
#                     uv-coordinate system within the plane across multiple TNs of a TJ.
#                     Under certain circumstances (when values of the normal vector cross
#                     zero), this can even lead to sudden 'flipping' of the orientation of
#                     the plane. Currently, this is implicitly being "fixed" downstream
#                     since the vector triplets are being aligned by rotation and flipping
#                     prior to their reduction to a consensus triplet. However, it might be
#                     more clean and robust to do something slightly more sophisticated than
#                     classical Gram-Schmidt in order to enforce consistency.

# Import sympy symbols
from sympy.abc import q,r,s,  x,y,z  # (normal vector), (point to be projected)

# Use Gram-Schmidt orthogonalization to create orthonormal vectors defining the in-plane
# coordinate system given three arbitrary vectors, the first of which is the normal vector
# of the plane. The other two (defining the in-plane directionalities) are arbitrarily
# chosen such that they will never fall onto the normal vector or onto each other.
orthonormals = sym.GramSchmidt([sym.Matrix([q,         r,         s]),  # Normal vec to plane -> first coordinate vec
                                sym.Matrix([q, 2*(r+0.1), 3*(s+0.1)]),  # Arbitrary vec not on the normal vec
                                sym.Matrix([2*(q+0.1), 3*(r+0.1), s])], # Arbitrary vec not on either other vec
                                orthonormal=True)           # Normalize resulting orthogonal vectors

# With the resulting orthonormals defining the new coordinate system, the projection
# of points into it is just a straightforward dot product.
projection_pt = sym.Matrix([x, y, z])
proj_d = orthonormals[0].dot(projection_pt)  # Distance from plane
proj_u = orthonormals[1].dot(projection_pt)  # Coordinate along first axis in plane
proj_v = orthonormals[2].dot(projection_pt)  # Coordinate along second axis in plane

# Lambdify
lambda_dist = sym.utilities.lambdify((q,r,s,x,y,z), proj_d, modules='numpy')
lambda_u    = sym.utilities.lambdify((q,r,s,x,y,z), proj_u, modules='numpy')
lambda_v    = sym.utilities.lambdify((q,r,s,x,y,z), proj_v, modules='numpy')

# Wrap into a function (sequential)
def p2p_projection(normal_vec, pt_coords):

    # Unpack inputs
    q,r,s = normal_vec[2], normal_vec[1], normal_vec[0]
    x,y,z = pt_coords[:,2], pt_coords[:,1], pt_coords[:,0]

    # Run projection
    dists = np.abs(lambda_dist(q,r,s,x,y,z))
    p_u   = lambda_u(q,r,s,x,y,z)
    p_v   = lambda_v(q,r,s,x,y,z)

    # Pack and return outputs
    projected = np.array([p_u, p_v]).T
    return projected, dists

## Wrap into a function (vectorized)
## FLAG -- ISSUE: This does not work as intended! It runs but does not yield the same results
##                as the sequential version. There is likely an missing `axis=` kwarg in one
##                of the numpy functions substituted by lambdify. This could perhaps be fixed
##                by manual inspection of the projection function.
#def p2p_projection_vectorized(normal_vec, pt_coords):
#
#    # Unpack inputs
#    q,r,s = normal_vec[..., 2, np.newaxis], normal_vec[..., 1, np.newaxis], normal_vec[..., 0, np.newaxis]
#    x,y,z = pt_coords[..., 2], pt_coords[..., 1], pt_coords[..., 0]
#
#    # Run projection
#    dists = np.abs(lambda_dist(q,r,s,x,y,z))
#    p_u   = lambda_u(q,r,s,x,y,z)
#    p_v   = lambda_v(q,r,s,x,y,z)
#
#    # Pack and return outputs
#    projected = np.rollaxis(np.array([p_u, p_v]), 2)
#    projected = np.rollaxis(projected, 2)
#    return projected, dists
```

```python
### Project relevant DNs onto the TJ-orthogonal plane

# Params
close_points_radius =  25.0
dist_points_keep    = 100

# Prep output dicts
TJs_DNs_proj = {}
TJs_DNs_dist = {}

# For each TJ...
for TJ_ID in TJs_spline_ev.keys():

    # Find the IDs of the three connected interfaces
    DJ_IDs = list(itertools.combinations(TJ_ID, 2))

    # Skip edge cases with more than 3
    if len(DJ_IDs) > 3:
        continue

    # Get corresponding TJ-normal vectors
    proj_tangents = TJs_spline_tangent[TJ_ID]

    # Prep output lists
    TJs_DNs_proj[TJ_ID] = {DJ_ID:[] for DJ_ID in DJ_IDs}
    TJs_DNs_dist[TJ_ID] = {DJ_ID:[] for DJ_ID in DJ_IDs}

    # For each TN of the current TJ...
    for TN_idx, TN in enumerate(TJs_spline_ev[TJ_ID]):

        # For each connected interface...
        for DJ_idx, DJ_ID in enumerate(DJ_IDs):

            # Get all the DJ points of that interface
            current_DNCs = DNCs[DJs[DJ_ID]]

            # If there are none, skip this TN
            if current_DNCs.size == 0:
                print("Skipped case at TJ_ID="+str(TJ_ID) + ", TN_idx=" +str(TN_idx) +
                      ", DJ_ID="+str(DJ_ID)+" ->> lacks interface points!")
                TJs_DNs_proj[TJ_ID][DJ_ID].append(np.empty(0))
                continue

            # Get the DN points close to the TN
            kdtree = spatial.cKDTree(current_DNCs)
            KNNs   = kdtree.query_ball_point(TN, close_points_radius)

            # If there are none, skip this TN
            if not KNNs:
                print("Skipped case at TJ_ID="+str(TJ_ID) + ", TN_idx=" +str(TN_idx) +
                      ", DJ_ID="+str(DJ_ID)+" ->> no close-by neighbors!")
                TJs_DNs_proj[TJ_ID][DJ_ID].append(np.empty(0))
                continue

            # Move the points onto the origin
            current_DNCs = current_DNCs[KNNs] - TJs_spline_ev[TJ_ID][TN_idx]

            # Nothing *should* go wrong here - but if it does, first look into
            # the way the arbitrary vectors for Gram-Schmidt are generated!
            with np.errstate(divide='raise', invalid='raise'):

                # Project the TN points onto the dihedral plane
                projs, dists = p2p_projection(proj_tangents[TN_idx], current_DNCs)

            # Threshold on the distances; keep at most the closest n points
            psort = np.argsort(dists)[:dist_points_keep]
            projs = projs[psort]
            dists = dists[psort]

            # Keep the results
            TJs_DNs_proj[TJ_ID][DJ_ID].append(projs)
            TJs_DNs_dist[TJ_ID][DJ_ID].append(dists)
```

```python
### Visualize the projections

@interact(TJ_ID=list(TJs_DNs_proj.keys()),
          TN_idx=(0,num_ts-1,1))
def plot_proj(TJ_ID=list(TJs_DNs_proj.keys())[0],
              TN_idx=num_ts//2):

    # Prep plot
    plt.figure(figsize=(6,6))

    # For each adjacent DJ...
    for DJ_ID in itertools.combinations(TJ_ID, 2):

        # Plot the projected points
        plt.scatter(TJs_DNs_proj[TJ_ID][DJ_ID][TN_idx][:,1],
                    TJs_DNs_proj[TJ_ID][DJ_ID][TN_idx][:,0],
                    c=TJs_DNs_dist[TJ_ID][DJ_ID][TN_idx],
                    cmap='viridis', alpha=0.5, lw=0)

    # Finish
    plt.xlabel('u')
    plt.ylabel('v')
    plt.tight_layout()
    plt.show()
```

### Retrieving Incident Vectors in the Dihedral Plane

**Note:** The arc fitting approach taken here is based on the second approach described in [this scipy cookbook entry](https://scipy-cookbook.readthedocs.io/items/Least_Squares_Circle.html). It could probably be further improved by using the third approach, i.e. by explicitly specifying the Jacobian function.

```python
### Functions for circular arc fitting

# FLAG -- PERFORMANCE: The arc fitting approach used here could be sped up by explicitly
#                      specifying a Jacobian function, see the markdown note above.

# Compute coordinates from angle
def circle(r, cx, cy, alpha):
    x = r*np.cos(alpha) + cx
    y = r*np.sin(alpha) + cy
    return np.array([y,x])

# Compute radius/radii given a center and a point/multiple points
def radius(xc, yc, x, y):
    return np.sqrt((x-xc)**2 + (y-yc)**2)

# Loss: distance of data points from mean circle
def circle_loss(c, x, y):
    radii = radius(c[0], c[1], x, y)
    return radii - radii.mean()

# Subtraction of n1 and n2, wrapping around at minimum and maximum
def wrap_sub(n1, n2, minimum=-np.pi, maximum=np.pi):
    s = n1 - n2
    try:
        s[s<=minimum] = maximum + (s[s<=minimum] - minimum)
        s[s>=maximum] = minimum + (s[s>=maximum] - maximum)
    except TypeError:
        if s <= minimum: s = maximum + (s - minimum)
        if s >= maximum: s = minimum + (s - maximum)
    return s
```

```python
### Find incident vectors for each TN based on circular arc fitting

# FLAG -- ROBUSTNESS: There is still an edge case in this where perfectly straight
#                     lines are fit with a completely wrong (very small) circle.
#                     Right now, this is handled as a "silly exception" for the
#                     synthetic test sample, where the middle line between cells is
#                     perfectly straight. The hope is that this will never occur
#                     in real data - but if it does, the curent handling will
#                     almost certainly fail, as it presupposes that the line is
#                     not only perfectly straight but also perfectly aligned with
#                     one of the image axes.

# Prep output dict
TJs_vec_proj = {}

# For each TJ...
for TJ_ID in TJs_DNs_proj.keys():

    # Prepare an appropriate result array
    TJs_vec_proj[TJ_ID] = np.empty((TJs_spline_ev[TJ_ID].shape[0], 3, 2))  # Num. of TNs, 3 vectors, 2 dimensions

    # For each adjacent DJ...
    for DJ_idx, DJ_ID in enumerate(list(itertools.combinations(TJ_ID, 2))):

        # For each TN along the TJ...
        for TN_idx in range(len(TJs_DNs_proj[TJ_ID][DJ_ID])):

            # Prep data for fitting
            x = TJs_DNs_proj[TJ_ID][DJ_ID][TN_idx][:,1]
            y = TJs_DNs_proj[TJ_ID][DJ_ID][TN_idx][:,0]

            # Catch silly exception where all the data is in a line (may cause artifacts)
            is_silly_exception = False
            if np.allclose(x, x[0]) or np.allclose(y, y[0]):
                is_silly_exception = True

            # Fit a circle to the data
            center, ier = optimize.leastsq(circle_loss, [np.mean(x), np.mean(y)], args=(x, y))
            cx, cy = center
            r      = radius(cx, cy, x, y).mean()

            # Get angular position of the TN point (which is the origin in the projection)
            TN_alpha = np.arctan2(0.0-cy, 0.0-cx)

            # Get correct sign for tangent vector direction
            DNs_alpha = wrap_sub(np.arctan2(y-cy, x-cx), TN_alpha)
            sign = np.sign(np.mean(DNs_alpha))

            # Get tangent vector based on TN angle and small shift
            TN_proj = circle(r, cx, cy, TN_alpha)
            shifted = circle(r, cx, cy, TN_alpha+10e-5)
            tangent = shifted - TN_proj
            tangent = tangent * sign

            # Handle the silly exception where all the data is in a line
            if is_silly_exception:
                tangent = np.array([np.mean(y), np.mean(x)])

            # Normalize to magnitude 1
            tangent = tangent / np.sqrt(np.sum(tangent**2.0))

            # Save the result
            TJs_vec_proj[TJ_ID][TN_idx, DJ_idx, :] = tangent
```

```python
### Visualize the projections

@interact(TJ_ID=list(TJs_DNs_proj.keys()),
          TN_idx=(0,num_ts-1,1))
def plot_proj(TJ_ID=list(TJs_DNs_proj.keys())[0],
              TN_idx=num_ts//2):

    # Prep plot
    plt.figure(figsize=(6,6))

    # For each adjacent DJ...
    for DJ_idx, DJ_ID in enumerate(list(itertools.combinations(TJ_ID, 2))):

        # Plot the projected points
        plt.scatter(TJs_DNs_proj[TJ_ID][DJ_ID][TN_idx][:,1],
                    TJs_DNs_proj[TJ_ID][DJ_ID][TN_idx][:,0],
                    c=TJs_DNs_dist[TJ_ID][DJ_ID][TN_idx],
                    cmap='viridis', alpha=0.5, lw=0)

        # Plot the fitted vectors
        plt.plot([0, TJs_vec_proj[TJ_ID][TN_idx, DJ_idx, 1]*10],
                 [0, TJs_vec_proj[TJ_ID][TN_idx, DJ_idx, 0]*10],
                 c='k', lw='2', alpha=0.75)

    # Finish
    plt.xlabel('u')
    plt.ylabel('v')
    plt.tight_layout()
    plt.show()
```

```python
### Show the resulting vector triplets

# FLAT: NOTE -- These are not aligned, which is okay but not ideal;
#               see ROBUSTNESS flag in sympy projection code.

# For each TJ...
cols = ['r','g','b']
for TJ_ID in TJs_vec_proj.keys():

    # Prep
    plt.figure()

    # Plot each vec...
    for vec in TJs_vec_proj[TJ_ID]:
        for i,v in enumerate(vec):
            plt.plot([0,v[1]], [0,v[0]], c=cols[i])

    # Finalize
    plt.title(str(TJ_ID))
    plt.xlabel('x'); plt.ylabel('y')
    plt.axis('equal')
    plt.tight_layout()
    plt.show()
```

### Aligning Incident Vectors Along TJs

```python
### Align triplets based on first vector & flip those that are the wrong way around

# For each TJ...
TJs_vec_aligned = {}
for TJ_ID in TJs_vec_proj.keys():

    # Prep output container
    triplets_aligned = np.empty_like(TJs_vec_proj[TJ_ID])

    ## Rotate each triplet to lay the first vector onto zero angle [sequential]
    #angles_zeroed = np.empty((TJs_vec_proj[TJ_ID].shape[0], 3))
    #for t,triplet in enumerate(TJs_vec_proj[TJ_ID]):
    #    angles_raw = np.arctan2(triplet[:,0], triplet[:,1])
    #    angles_zeroed[t] = angles_raw - angles_raw[0]

    # Rotate each triplet to lay the first vector onto zero angle [vectorized]
    angles_raw = np.arctan2(TJs_vec_proj[TJ_ID][:,:,0], TJs_vec_proj[TJ_ID][:,:,1])
    angles_zeroed = angles_raw - angles_raw[:, 0, np.newaxis]

    ## Function: if flipped is better than the consensus (here the median), then flip [sequential]
    #def flip_improvement(angles_zeroed):
    #    median = np.median(angles_zeroed, axis=0)
    #    for t in range(len(angles_zeroed)):
    #        diff_original = np.abs(wrap_sub( angles_zeroed[t], median))
    #        diff_flipped  = np.abs(wrap_sub(-angles_zeroed[t], median))
    #        if np.sum(diff_flipped) < np.sum(diff_original):
    #            angles_zeroed[t] = - angles_zeroed[t]
    #    return angles_zeroed

    # Function: if flipped is better than the consensus (here the median), then flip [vectorized]
    def flip_improvement(angles_zeroed):
        median = np.median(angles_zeroed, axis=0)
        diff_original = np.abs(wrap_sub( angles_zeroed, median))
        diff_flipped  = np.abs(wrap_sub(-angles_zeroed, median))
        flip_mask = np.sum(diff_flipped, axis=1) < np.sum(diff_original, axis=1)
        angles_zeroed[flip_mask] = -angles_zeroed[flip_mask]
        return angles_zeroed

    # Run flip improvement until there is either...
    # ...no change from one step to the next, or
    # ...no improvement since 5 steps ago
    median  = np.median(angles_zeroed, axis=0)
    losses  = [np.abs(wrap_sub(angles_zeroed, median))]
    counter = 0
    while True:

        # Run a flip
        angles_zeroed_new = flip_improvement(angles_zeroed)

        # Break if it changed nothing
        if np.all(angles_zeroed==angles_zeroed_new):
            break

        # Otherwise, compute and keep the new loss
        median = np.median(angles_zeroed_new, axis=0)
        losses.append(np.abs(wrap_sub(angles_zeroed_new, median)))

        # Break if the new loss is worse or equal to the loss 5 steps ago
        if (counter >= 5) and losses[-1] >= losses[-6]:
            break

        # Update
        angles_zeroed = angles_zeroed_new
        counter += 1

    # Convert back to unit vectors
    for t in range(len(angles_zeroed)):
        triplets_aligned[t] = circle(1.0, 0.0, 0.0, angles_zeroed[t]).T

    # Store results
    TJs_vec_aligned[TJ_ID] = triplets_aligned
```

```python
### Show the resulting aligned vector triplets

# For each TJ...
cols = ['r','g','b']
for TJ_ID in TJs_vec_aligned.keys():

    # Prep
    plt.figure()

    # Plot each vec...
    for vec in TJs_vec_aligned[TJ_ID]:
        for i,v in enumerate(vec):
            plt.plot([0,v[1]], [0,v[0]], c=cols[i])

    # Finalize
    plt.title(str(TJ_ID))
    plt.xlabel('v'); plt.ylabel('u')
    plt.axis('equal')
    plt.tight_layout()
    plt.show()
```

### Finding Consensus Incident Vector Triplets

```python
### Generate a consensus incident vector triplet for each TJ

# FLAG -- PRECISION, FLAG -- ROBUSTNESS: This is currently done in a very simple fashion.
#                                        It probably works fine / doesn't matter much for
#                                        data with a high z-resolution. However, there is
#                                        room for improvement for low z-resolution data!

# For each TJ...
TJs_vec_consensus = {}
for TJ_ID in TJs_vec_aligned.keys():

    # Prep result container
    TJs_vec_consensus[TJ_ID] = np.empty((3,2))

    # Compute the mean of vectors within the 25-75th percentile
    p25, p75 = np.percentile(TJs_vec_aligned[TJ_ID], [25, 75], axis=0)
    for v in range(3):
        vec  = TJs_vec_aligned[TJ_ID][:,v,:]
        mask = (vec >= p25[v]) & (vec <= p75[v])
        mean_v = np.mean(vec[mask[:,0],0])
        mean_u = np.mean(vec[mask[:,1],1])
        TJs_vec_consensus[TJ_ID][v] = [mean_v, mean_u]
```

```python
### Show the resulting consensus vector triplets

# For each TJ...
cols = ['r','g','b']
for TJ_ID in TJs_vec_consensus.keys():

    # Prep
    plt.figure()

    # Plot individual vecs...
    for vec in TJs_vec_aligned[TJ_ID]:
        for i,v in enumerate(vec):
            plt.plot([0,v[1]], [0,v[0]], c=cols[i], alpha=0.1)

    # Plot consensus vecs
    for i,v in enumerate(TJs_vec_consensus[TJ_ID]):
        plt.plot([0,v[1]], [0,v[0]], c=cols[i], lw=4)

    # Finalize
    plt.title(str(TJ_ID))
    plt.xlabel('v'); plt.ylabel('u')
    plt.axis('equal')
    plt.tight_layout()
    plt.show()
```

### Solving the Force Balance Equations

#### Prep: Assembling Equation Matrix G

```python
### Assemble G

# Initialize zero matrix of shape (2 * num of TJs, num of DJs)
G = np.zeros((2*len(TJs_vec_consensus), len(DJs)))

# For each TJ...
for TJ_idx, TJ_ID in enumerate(TJs_vec_consensus.keys()):

    # Get all relevant DJs
    DJ_IDs  = list(itertools.combinations(TJ_ID, 2))

    # For each DJ...
    for DJ_ref, DJ_ID in enumerate(DJ_IDs):

        # Get index (in G) of the current DJ
        DJ_idx = list(DJs.keys()).index(DJ_ID)

        # Fill the appropriate positions in G
        G[TJ_idx, DJ_idx] = TJs_vec_consensus[TJ_ID][DJ_ref][0]
        G[len(TJs_vec_consensus)+TJ_idx, DJ_idx] = TJs_vec_consensus[TJ_ID][DJ_ref][1]
```

```python
### Visualize the result
plt.imshow(G)
plt.show()
```

#### Solve using the `contraints` kwarg of scipy's minimize

```python
### Define loss and constraints

# Loss: sum of square deviations of equilibrium equations
def eq_loss_c(gammas, G):
    loss = (np.dot(G, gammas))**2.0
    return np.sum(loss)

# Constraint: the mean of tensions must be 1
def eq_constraint(gammas):
    c = np.mean(gammas) - 1
    return c
```

```python
### Run the fit
fit = optimize.minimize(eq_loss_c, np.ones(len(DJs)), args=(G,),
                        constraints={'type':'eq', 'fun':eq_constraint})
tensions_c = fit.x
print(tensions_c)

# FLAG -- ISSUE: The softest interface has a negative tension. I'm not sure
#                if that indicates a problem. It might be perfectly fine;
#                after all, the tensions are relative to the mean and they
#                are effective surface tensions, so high adhesion should
#                be able to make them net-negative. Furthermore, this
#                synthetic test sample has been constructed from geometric
#                objects, so it doesn't represent a realistic structure.
```

```python
### Show tensions on image stack

@interact(z=(0, im.shape[0]-1, 1))
def show_stack(z=im.shape[0]//2):

    # Prep and plot image
    plt.figure(figsize=(8,8))
    plt.imshow(outlines_id[z], cmap='gray')

    # For each DJ...
    for DJ_num, DJ_ID in enumerate(DJs.keys()):

        # Get the DJ's DNs in the selected z plane
        DNs_in_plane = DNIs[DJs[DJ_ID]][DNIs[DJs[DJ_ID]][:,0]==z]

        # Plot the points
        plt.scatter(DNs_in_plane[:, 2], DNs_in_plane[:, 1],
                    c=[tensions_c[DJ_num] for _ in range(DNs_in_plane.shape[0])],
                    vmin=np.min(tensions_c), vmax=np.max(tensions_c),
                    cmap='viridis', s=20)

    # Finish
    plt.colorbar()
    plt.show()
```

#### Solve using Brodland et alii's Lagrange Multiplier Approach

```python
### Prepare the matrix

# Gdot
Gdot = np.dot(G.T, G)

# Show
plt.imshow(Gdot)
plt.show()

# Add the constraints
Gready = np.zeros((Gdot.shape[0]+1, Gdot.shape[1]+1))
Gready[:Gdot.shape[0], :Gdot.shape[1]] = Gdot
Gready[-1,:-1] = 1.0
Gready[:-1,-1] = 1.0

# Show
plt.imshow(Gready)
plt.show()
```

```python
### Define loss

# Loss combining tension fit and constraint
def eq_loss_l(gammas_lagrange, Gready):
    loss = np.sum(np.dot(Gready[:-1], gammas_lagrange)**2.0)  # Fit loss
    loss += (np.dot(Gready[-1], gammas_lagrange) - (gammas_lagrange.size-1))**2.0  # Constraint loss
    return loss
```

```python
### Run the fit
fit = optimize.minimize(eq_loss_l, np.ones(len(DJs)+1), args=(Gready,))
tensions_l = fit.x[:-1]
lagrange = fit.x[-1]
print(tensions_l)
print(lagrange)

# FLAG -- ISSUE: Negative tension value, same as above with the scipy-based
#                approach. See flag there for more info.
```

```python
### Show tensions on image stack

@interact(z=(0, im.shape[0]-1, 1))
def show_stack(z=im.shape[0]//2):

    # Prep and plot image
    plt.figure(figsize=(8,8))
    plt.imshow(outlines_id[z], cmap='gray')

    # For each DJ...
    for DJ_num, DJ_ID in enumerate(DJs.keys()):

        # Get the DJ's DNs in the selected z plane
        DNs_in_plane = DNIs[DJs[DJ_ID]][DNIs[DJs[DJ_ID]][:,0]==z]

        # Plot the points
        plt.scatter(DNs_in_plane[:, 2], DNs_in_plane[:, 1],
                    c=[tensions_l[DJ_num] for _ in range(DNs_in_plane.shape[0])],
                    vmin=np.min(tensions_l), vmax=np.max(tensions_l),
                    cmap='viridis', s=20)

    # Finish
    plt.colorbar()
    plt.show()
```

#### Compare Solver vs Lagrange

```python
### Plot against each other

# Prep
plt.figure(figsize=(5,5))

# Plot
plt.scatter(tensions_l, tensions_c, s=50,
            c='darkblue', lw=0.5, edgecolor='cyan')

# Add equality line
xlims, ylims = plt.gca().get_xlim(), plt.gca().get_ylim()
plt.plot([-10,10], [-10,10], 'k-', zorder=-1, lw=1, alpha=0.5)
plt.xlim(xlims); plt.ylim(ylims)

# Labels
plt.xlabel("inferred tension\n[lagrange multiplier]")
plt.ylabel("inferred tension\n[scipy constraint]")

# Finalize
plt.tight_layout()
plt.show()
```


```python
%matplotlib inline
```

```python
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sn
sn.set_style("whitegrid")
import sympy as sym

import pycollocation
import pyam
```

```python
# should be 0.4.0a0
pycollocation.__version__
```

    '0.4.0a0'

```python
# should be 0.2.0a0
pyam.__version__
```

    '0.2.0a0'

## Defining inputs

Need to define some heterogenous factors of production...

```python
# define some workers skill
x, loc1, mu1, sigma1 = sym.var('x, loc1, mu1, sigma1')
skill_cdf = 0.5 + 0.5 * sym.erf((sym.log(x - loc1) - mu1) / sym.sqrt(2 * sigma1**2))
skill_params = {'loc1': 1e0, 'mu1': 0.0, 'sigma1': 1.0}

workers = pyam.Input(var=x,
                     cdf=skill_cdf,
                     params=skill_params,
                     bounds=(1.2, 1e1),  # guesses for the alpha and (1 - alpha) quantiles!
                     alpha=0.005,
                     measure=15.0  # 15x more workers than firms
                     )

# define some firms
y, loc2, mu2, sigma2 = sym.var('y, loc2, mu2, sigma2')
productivity_cdf = 0.5 + 0.5 * sym.erf((sym.log(y - loc2) - mu2) / sym.sqrt(2 * sigma2**2))
productivity_params = {'loc2': 1e0, 'mu2': 0.0, 'sigma2': 1.0}

firms = pyam.Input(var=y,
                   cdf=productivity_cdf,
                   params=productivity_params,
                   bounds=(1.2, 1e1),  # guesses for the alpha and (1 - alpha) quantiles!
                   alpha=0.005,
                   measure=1.0
                   )
```

Note that we are shifting the distributions of worker skill and firm productivity to the right by 1.0 in order to try and avoid issues with having workers (firms) with near zero skill (productivity).

```python
xs = np.linspace(workers.lower, workers.upper, 1e4)
plt.plot(xs, workers.evaluate_pdf(xs))
plt.xlabel('Worker skill, $x$', fontsize=20)
plt.show()
```

## Defining a production process

Next need to define some production process...

```python
# define symbolic expression for CES between x and y
omega_A, sigma_A = sym.var('omega_A, sigma_A')
A = ((omega_A * x**((sigma_A - 1) / sigma_A) +
     (1 - omega_A) * y**((sigma_A - 1) / sigma_A))**(sigma_A / (sigma_A - 1)))

# define symbolic expression for CES between x and y
r, l, omega_B, sigma_B = sym.var('r, l, omega_B, sigma_B')
B = l**omega_B * r**(1 - omega_B)
F = A * B
```

```python
# positive assortativity requires that sigma_A * sigma_B < 1
F_params = {'omega_A':0.25, 'omega_B':0.5, 'sigma_A':0.5, 'sigma_B':1.0 }
```

## Define a boundary value problem

```python
problem = pyam.AssortativeMatchingProblem(assortativity='positive',
                                          input1=workers,
                                          input2=firms,
                                          F=F,
                                          F_params=F_params)
```

## Pick some collocation solver

```python
solver = pycollocation.OrthogonalPolynomialSolver(problem)
```

## Compute some decent initial guess
Currently I guess that $\mu(x)$ is has the form...

$$ \hat{\mu}(x) = \beta_0 + \beta_1 f(x) $$

(i.e., a linear translation) of some function $f$.  Using my $\hat{\mu}(x)$, I can then back out a guess for $\theta(x)$ implied by the model...

$$ \hat{\theta}(x) = \frac{H(x)}{\hat{\mu}'(x)} $$

```python
initial_guess = pyam.OrthogonalPolynomialInitialGuess(solver)
initial_polys = initial_guess.compute_initial_guess("Chebyshev",
                                                    degrees={'mu': 75, 'theta': 75},
                                                    f=lambda x, alpha: x**alpha,
                                                    alpha=0.25)
```

```python
# quickly plot the initial conditions
xs = np.linspace(workers.lower, workers.upper, 1000)
plt.plot(xs, initial_polys['mu'](xs))
plt.plot(xs, initial_polys['theta'](xs))
plt.grid('on')
```

## Solve the model!

```python
domain = [workers.lower, workers.upper]
initial_coefs = {'mu': initial_polys['mu'].coef,
                 'theta': initial_polys['theta'].coef}

solver.solve(kind="Chebyshev",
             coefs_dict=initial_coefs,
             domain=domain,
             method='hybr')
```

```python
solver.result.success
```

    True

## Plot some results

```python
viz = pyam.Visualizer(solver)
```

```python
viz.interpolation_knots = np.linspace(workers.lower, workers.upper, 1000)
viz.residuals.plot()
plt.show()
```

```python
viz.normalized_residuals[['mu', 'theta']].plot(logy=True)
plt.show()
```

```python
viz.solution.tail()
```

      F
      Fl
      Flr
      Fr
      Fx
      Fxl
      Fxr
      Fxy
      Fy
      Fyl
      Fyr
      factor_payment_1
      factor_payment_2
      mu
      theta

      14.089895
      80.393180
      1.239461
      0.619731
      40.196590
      1.429173
      0.022034
      0.714586
      0.151663
      4.265630
      0.065765
      2.132815
      1.239461
      40.196590
      14.125999
      32.430700

      14.102974
      80.427887
      1.240037
      0.620019
      40.213944
      1.427779
      0.022013
      0.713889
      0.151495
      4.266927
      0.065787
      2.133463
      1.240037
      40.213944
      14.130067
      32.429627

      14.116053
      80.462524
      1.240613
      0.620306
      40.231262
      1.426386
      0.021993
      0.713193
      0.151328
      4.268222
      0.065810
      2.134111
      1.240613
      40.231262
      14.134125
      32.428541

      14.129132
      80.497089
      1.241188
      0.620594
      40.248544
      1.424994
      0.021972
      0.712497
      0.151162
      4.269516
      0.065832
      2.134758
      1.241188
      40.248544
      14.138173
      32.427442

      14.142212
      80.531577
      1.241762
      0.620881
      40.265788
      1.423603
      0.021951
      0.711801
      0.150995
      4.270809
      0.065854
      2.135404
      1.241762
      40.265788
      14.142212
      32.426326

```python
viz.solution[['mu', 'theta']].plot(subplots=True)
plt.show()
```

```python
viz.solution[['Fxy', 'Fyl']].plot()
plt.show()
```

## Plot factor payments
Note the `factor_payment_1` is wages and `factor_payment_2` is profits...

```python
viz.solution[['factor_payment_1', 'factor_payment_2']].plot(subplots=True)
plt.show()
```

## Plot firm size against wages and profits

```python
fig, axes = plt.subplots(1, 2, sharey=True)
axes[0].scatter(viz.solution.factor_payment_1, viz.solution.theta, alpha=0.5,
               edgecolor='none')
axes[0].set_ylim(0, 1.05 * viz.solution.theta.max())
axes[0].set_xlabel('Wages, $w$')
axes[0].set_ylabel(r'Firm size, $\theta$')

axes[1].scatter(viz.solution.factor_payment_2, viz.solution.theta, alpha=0.5,
               edgecolor='none')
axes[1].set_xlabel(r'Profits, $\pi$')

plt.show()
```

```python
# to get correlation just use pandas!
viz.solution.corr()
```

      F
      Fl
      Flr
      Fr
      Fx
      Fxl
      Fxr
      Fxy
      Fy
      Fyl
      Fyr
      factor_payment_1
      factor_payment_2
      mu
      theta

      F
      1.000000
      0.999021
      0.999021
      1.000000
      0.220809
      -0.997588
      0.220809
      -0.972950
      0.992603
      -0.614926
      0.992603
      0.999021
      1.000000
      0.995928
      0.946557

      Fl
      0.999021
      1.000000
      1.000000
      0.999021
      0.235158
      -0.999270
      0.235158
      -0.976912
      0.996132
      -0.634471
      0.996132
      1.000000
      0.999021
      0.996189
      0.950320

      Flr
      0.999021
      1.000000
      1.000000
      0.999021
      0.235158
      -0.999270
      0.235158
      -0.976912
      0.996132
      -0.634471
      0.996132
      1.000000
      0.999021
      0.996189
      0.950320

      Fr
      1.000000
      0.999021
      0.999021
      1.000000
      0.220809
      -0.997588
      0.220809
      -0.972950
      0.992603
      -0.614926
      0.992603
      0.999021
      1.000000
      0.995928
      0.946557

      Fx
      0.220809
      0.235158
      0.235158
      0.220809
      1.000000
      -0.271117
      1.000000
      -0.434857
      0.309551
      -0.842063
      0.309551
      0.235158
      0.220809
      0.307716
      0.523544

      Fxl
      -0.997588
      -0.999270
      -0.999270
      -0.997588
      -0.271117
      1.000000
      -0.271117
      0.984288
      -0.998293
      0.659932
      -0.998293
      -0.999270
      -0.997588
      -0.998115
      -0.960972

      Fxr
      0.220809
      0.235158
      0.235158
      0.220809
      1.000000
      -0.271117
      1.000000
      -0.434857
      0.309551
      -0.842063
      0.309551
      0.235158
      0.220809
      0.307716
      0.523544

      Fxy
      -0.972950
      -0.976912
      -0.976912
      -0.972950
      -0.434857
      0.984288
      -0.434857
      1.000000
      -0.988766
      0.763174
      -0.988766
      -0.976912
      -0.972950
      -0.989287
      -0.993711

      Fy
      0.992603
      0.996132
      0.996132
      0.992603
      0.309551
      -0.998293
      0.309551
      -0.988766
      1.000000
      -0.699246
      1.000000
      0.996132
      0.992603
      0.996621
      0.969120

      Fyl
      -0.614926
      -0.634471
      -0.634471
      -0.614926
      -0.842063
      0.659932
      -0.842063
      0.763174
      -0.699246
      1.000000
      -0.699246
      -0.634471
      -0.614926
      -0.676287
      -0.814078

      Fyr
      0.992603
      0.996132
      0.996132
      0.992603
      0.309551
      -0.998293
      0.309551
      -0.988766
      1.000000
      -0.699246
      1.000000
      0.996132
      0.992603
      0.996621
      0.969120

      factor_payment_1
      0.999021
      1.000000
      1.000000
      0.999021
      0.235158
      -0.999270
      0.235158
      -0.976912
      0.996132
      -0.634471
      0.996132
      1.000000
      0.999021
      0.996189
      0.950320

      factor_payment_2
      1.000000
      0.999021
      0.999021
      1.000000
      0.220809
      -0.997588
      0.220809
      -0.972950
      0.992603
      -0.614926
      0.992603
      0.999021
      1.000000
      0.995928
      0.946557

      mu
      0.995928
      0.996189
      0.996189
      0.995928
      0.307716
      -0.998115
      0.307716
      -0.989287
      0.996621
      -0.676287
      0.996621
      0.996189
      0.995928
      1.000000
      0.971759

      theta
      0.946557
      0.950320
      0.950320
      0.946557
      0.523544
      -0.960972
      0.523544
      -0.993711
      0.969120
      -0.814078
      0.969120
      0.950320
      0.946557
      0.971759
      1.000000

```python
# or a subset
viz.solution[['theta', 'factor_payment_1']].corr()
```

      theta
      factor_payment_1

      theta
      1.00000
      0.95032

      factor_payment_1
      0.95032
      1.00000

```python
# or actual values!
viz.solution.corr().loc['theta']['factor_payment_1']
```

    0.95031954348391101

## Plot the density for firm size

As you can see, the theta function is hump-shaped. Nothing special, but when calculating the pdf some arrangements have to be done for this: sort the thetas preserving the order (so we can relate them to their xs) and then use carefully the right x for calculating the pdf.

The principle of Philipp's trick is:

$pdf_x(x_i)$ can be interpreted as *number of workers with ability x*. $\theta_i$ is the size of the firms that employs workers of kind $x_i$. As all firms that match with workers type $x_i$ choose the same firm size, $pdf_x(x_i)/\theta_i$ is the number of firms of size $\theta_i$.

Say there are 100 workers with ability $x_i$, and their associated firm size $\theta_i$ is 2. Then there are $100/2 = 50$ $ \theta_i$ firms

```python
fig, axes = plt.subplots(1, 3)
theta_pdf = viz.compute_pdf('theta', normalize=True)
theta_pdf.plot(ax=axes[0])
axes[0].set_xlabel(r'Firm size, $\theta$')
axes[0].set_title(r'pdf')

theta_cdf = viz.compute_cdf(theta_pdf)
theta_cdf.plot(ax=axes[1])
axes[1].set_title(r'cdf')
axes[1].set_xlabel(r'Firm size, $\theta$')

theta_sf = viz.compute_sf(theta_cdf)
theta_sf.plot(ax=axes[2])
axes[2].set_title(r'sf')
axes[2].set_xlabel(r'Firm size, $\theta$')

plt.tight_layout()
plt.show()
```

## Distributions of factor payments

Can plot the distributions of average factor payments...

```python
fig, axes = plt.subplots(1, 3)
factor_payment_1_pdf = viz.compute_pdf('factor_payment_1', normalize=True)
factor_payment_1_pdf.plot(ax=axes[0])
axes[0].set_title(r'pdf')

factor_payment_1_cdf = viz.compute_cdf(factor_payment_1_pdf)
factor_payment_1_cdf.plot(ax=axes[1])
axes[1].set_title(r'cdf')

factor_payment_1_sf = viz.compute_sf(factor_payment_1_cdf)
factor_payment_1_sf.plot(ax=axes[2])
axes[2].set_title(r'sf')

plt.tight_layout()
plt.show()
```

```python
fig, axes = plt.subplots(1, 3)
factor_payment_2_pdf = viz.compute_pdf('factor_payment_2', normalize=True)
factor_payment_2_pdf.plot(ax=axes[0])
axes[0].set_title(r'pdf')

factor_payment_2_cdf = viz.compute_cdf(factor_payment_2_pdf)
factor_payment_2_cdf.plot(ax=axes[1])
axes[1].set_title(r'cdf')

factor_payment_2_sf = viz.compute_sf(factor_payment_2_cdf)
factor_payment_2_sf.plot(ax=axes[2])
axes[2].set_title(r'sf')

plt.tight_layout()
plt.show()
```

## Widget

```python
from IPython.html import widgets
```

```python
def interactive_plot(viz, omega_A=0.25, omega_B=0.5, sigma_A=0.5, sigma_B=1.0,
                     loc1=1.0, mu1=0.0, sigma1=1.0, loc2=1.0, mu2=0.0, sigma2=1.0):
    # update new parameters as needed
    new_F_params = {'omega_A': omega_A, 'omega_B': omega_B,
                    'sigma_A': sigma_A, 'sigma_B': sigma_B}
    viz.solver.problem.F_params = new_F_params

    new_input1_params = {'loc1': loc1, 'mu1': mu1, 'sigma1': sigma1}
    viz.solver.problem.input1.params = new_input1_params

    new_input2_params = {'loc2': loc2, 'mu2': mu2, 'sigma2': sigma2}
    viz.solver.problem.input2.params = new_input2_params

    # solve the model using a hotstart initial guess
    domain = [viz.solver.problem.input1.lower, viz.solver.problem.input1.upper]
    initial_coefs = viz.solver._coefs_array_to_dict(viz.solver.result.x, viz.solver.degrees)
    viz.solver.solve(kind="Chebyshev",
                     coefs_dict=initial_coefs,
                     domain=domain,
                     method='hybr')

    if viz.solver.result.success:
        viz._Visualizer__solution = None  # should not need to access this!
        viz.interpolation_knots = np.linspace(domain[0], domain[1], 1000)
        viz.solution[['mu', 'theta']].plot(subplots=True)
        viz.normalized_residuals[['mu', 'theta']].plot(logy=True)
    else:
        print("Foobar!")
```

```python
viz_widget = widgets.fixed(viz)

# widgets for the model parameters
eps = 1e-2
omega_A_widget = widgets.FloatSlider(value=0.25, min=eps, max=1-eps, step=eps,
                                     description=r"$\omega_A$")
sigma_A_widget = widgets.FloatSlider(value=0.5, min=eps, max=1-eps, step=eps,
                                     description=r"$\sigma_A$")
omega_B_widget = widgets.FloatSlider(value=0.5, min=eps, max=1-eps, step=eps,
                                     description=r"$\omega_B$")
sigma_B_widget = widgets.fixed(1.0)

# widgets for input distributions
loc_widget = widgets.fixed(1.0)
mu_1_widget = widgets.FloatSlider(value=0.0, min=-1.0, max=1.0, step=eps,
                                  description=r"$\mu_1$")
mu_2_widget = widgets.FloatSlider(value=0.0, min=-1.0, max=1.0, step=eps,
                                  description=r"$\mu_2$")
sigma_1_widget = widgets.FloatSlider(value=1.0, min=eps, max=2-eps, step=eps,
                                     description=r"$\sigma_1$")
sigma_2_widget = widgets.FloatSlider(value=1.0, min=eps, max=2-eps, step=eps,
                                     description=r"$\sigma_2$")

```

```python
widgets.interact(interactive_plot, viz=viz_widget, omega_A=omega_A_widget,
                 sigma_A=sigma_A_widget, omega_B=omega_B_widget,
                 sigma_B=sigma_B_widget, sigma1=sigma_1_widget,
                 loc1=loc_widget, mu1 = mu_1_widget,
                 loc2=loc_widget, sigma2=sigma_2_widget, mu2 = mu_2_widget)
```

```python
# widget is changing the parameters of the underlying solver
solver.result.x
```

    array([  8.90568813e+00,   6.50304514e+00,  -1.34672823e+00,
             4.28680104e-02,   4.77357817e-02,  -1.16819661e-02,
             2.42768746e-03,  -2.37437996e-03,   2.10810240e-03,
            -1.40176909e-03,   7.99162058e-04,  -3.80275138e-04,
             9.64930130e-05,   9.20733908e-05,  -2.11980720e-04,
             2.83286219e-04,  -3.21237323e-04,   3.36934880e-04,
            -3.38214224e-04,   3.30504837e-04,  -3.17512949e-04,
             3.01729493e-04,  -2.84799791e-04,   2.67787975e-04,
            -2.51363021e-04,   2.35927790e-04,  -2.21707416e-04,
             2.08809016e-04,  -1.97261326e-04,   1.87040266e-04,
            -1.78084533e-04,   1.70303868e-04,  -1.63581505e-04,
             1.57771621e-04,  -1.52692795e-04,   1.48115337e-04,
            -1.43718237e-04,   1.38974317e-04,  -1.33162529e-04,
             1.26395135e-04,  -6.07009143e-05,   2.41292593e+00,
             1.26940051e+00,  -5.48938116e-01,   1.31155224e-01,
            -2.84148111e-02,   8.73652285e-03,  -2.25182023e-03,
            -5.50594141e-04,   1.27830608e-03,  -1.14200127e-03,
             7.96414650e-04,  -4.70396075e-04,   2.19107194e-04,
            -4.34628079e-05,  -7.07512035e-05,   1.39940506e-04,
            -1.78159629e-04,   1.96029790e-04,  -2.01052117e-04,
             1.98306933e-04,  -1.91133904e-04,   1.81675550e-04,
            -1.71275286e-04,   1.60754670e-04,  -1.50600201e-04,
             1.41086195e-04,  -1.32354122e-04,   1.24462970e-04,
            -1.17420636e-04,   1.11203010e-04,  -1.05765079e-04,
             1.01046828e-04,  -9.69756361e-05,   9.34662953e-05,
            -9.04198341e-05,   8.77217768e-05,  -8.52328118e-05,
             8.27571092e-05,  -8.00712163e-05,   7.73519043e-05,
            -3.75586336e-05])

```python

```