如何计算成对距离函数的雅可比行列式 (`scipy.spatial.pdist`)

Question

语境

我是 netgraph 的作者和维护者，这是一个用于创建网络可视化的 python 库。 我目前正在尝试优化一个例程，该例程计算一组 N 网络的节点位置，其中每条边都具有定义的长度。 可以找到一个例子here。

问题

该例程的核心是运行scipy.optimize.minimize 来计算使节点之间的总距离最大化的位置：

def cost_function(positions):
    return 1. / np.sum((pdist(positions.reshape((-1, 2))))**power)

result = minimize(cost_function, initial_positions.flatten(), method='SLSQP',
                  jac="2-point", constraints=[nonlinear_constraint])

• positions 是 (unravelled) (x, y) 元组的 numpy 数组。
• power 是一个较小的数字，它限制了大距离的影响（以鼓励紧凑的节点布局），但出于这个问题的目的，可以假定为 1。
• pdist是scipy.spatial中的成对距离函数。

使用以下非线性约束来约束最小化（/最大化）：

lower_bounds = ... # (squareform of an) (N, N) distance matrix of the sum of node sizes (i.e. nodes should not overlap)
upper_bounds = ... # (squareform of an) (N, N) distance matrix constructed from the given edge lengths

def constraint_function(positions):
    positions = np.reshape(positions, (-1, 2))
    return pdist(positions)

nonlinear_constraint = NonlinearConstraint(constraint_function, lb=lower_bounds, ub=upper_bounds, jac='2-point')

对于玩具示例，优化可以正确快速地完成。然而，即使对于小型网络，运行时间也相当糟糕。 我当前的实现使用有限差分来近似梯度 (jac='2-point')。 为了加快计算速度，我想明确地计算雅可比矩阵。

在几个 Math Stackexchange 帖子（1、2）之后，我计算了成对距离函数的雅可比矩阵，如下所示：

    def delta_constraint(positions):
        positions = np.reshape(positions, (-1, 2))
        total_positions = positions.shape[0]
        delta = positions[np.newaxis, :, :] - positions[:, np.newaxis, :]
        distance = np.sqrt(np.sum(delta ** 2, axis=-1))
        jac = delta / distance[:, :, np.newaxis]
        squareform_indices = np.triu_indices(total_positions, 1)
        return jac[squareform_indices]

nonlinear_constraint = NonlinearConstraint(constraint_function, lb=lower_bounds, ub=upper_bounds, jac=delta_constraint)

但是，这会导致 ValueError，因为输出的形状不正确。对于三角形示例，预期的输出形状是 (3, 6)，而上面的函数返回一个 (3, 2) 数组（即 3 个成对距离乘以 2 个维度）。对于正方形，预期输出为 (6, 8)，而实际输出为 (6, 2)。对于为 NonlinearConstraint 和 minimize 的 jac 参数实现正确的可调用项的任何帮助将不胜感激。

笔记

我想避免使用 autograd/jax/numdifftools（如 this question），因为我想使我的库的依赖项数量较少。

最小的工作示例

#!/usr/bin/env python
"""
Create a node layout with fixed edge lengths but unknown node positions.
"""

import numpy as np

from scipy.optimize import minimize, NonlinearConstraint
from scipy.spatial.distance import pdist, squareform


def get_geometric_node_layout(edges, edge_length, node_size=0., power=0.2, maximum_iterations=200, origin=(0, 0), scale=(1, 1)):
    """Node layout for defined edge lengths but unknown node positions.

    Node positions are determined through non-linear optimisation: the
    total distance between nodes is maximised subject to the constraint
    imposed by the edge lengths, which are used as upper bounds.
    If provided, node sizes are used to set lower bounds.

    Parameters
    ----------
    edges : list
        The edges of the graph, with each edge being represented by a (source node ID, target node ID) tuple.
    edge_lengths : dict
        Mapping of edges to their lengths.
    node_size : scalar or dict, default 0.
        Size (radius) of nodes.
        Providing the correct node size minimises the overlap of nodes in the graph,
        which can otherwise occur if there are many nodes, or if the nodes differ considerably in size.
    power : float, default 0.2.
        The cost being minimised is the inverse of the sum of distances.
        The power parameter is the exponent applied to each distance before summation.
        Large values result in positions that are stretched along one axis.
        Small values decrease the influence of long distances on the cost
        and promote a more compact layout.
    maximum_iterations : int
        Maximum number of iterations of the minimisation.
    origin : tuple, default (0, 0)
        The (float x, float y) coordinates corresponding to the lower left hand corner of the bounding box specifying the extent of the canvas.
    scale : tuple, default (1, 1)
        The (float x, float y) dimensions representing the width and height of the bounding box specifying the extent of the canvas.

    Returns
    -------
    node_positions : dict
        Dictionary mapping each node ID to (float x, float y) tuple, the node position.

    """
    # TODO: assert triangle inequality

    # TODO: assert that the edges fit within the canvas dimensions

    # ensure that graph is bi-directional
    edges = edges + [(target, source) for (source, target) in edges] # forces copy
    edges = list(set(edges))

    # upper bound: pairwise distance matrix with unknown distances set to the maximum possible distance given the canvas dimensions
    lengths = []
    for (source, target) in edges:
        if (source, target) in edge_length:
            lengths.append(edge_length[(source, target)])
        else:
            lengths.append(edge_length[(target, source)])

    sources, targets = zip(*edges)
    nodes = sources + targets
    unique_nodes = set(nodes)
    indices = range(len(unique_nodes))
    node_to_idx = dict(zip(unique_nodes, indices))
    source_indices = [node_to_idx[source] for source in sources]
    target_indices = [node_to_idx[target] for target in targets]

    total_nodes = len(unique_nodes)
    max_distance = np.sqrt(scale[0]**2 + scale[1]**2)
    distance_matrix = np.full((total_nodes, total_nodes), max_distance)
    distance_matrix[source_indices, target_indices] = lengths
    distance_matrix[np.diag_indices(total_nodes)] = 0
    upper_bounds = squareform(distance_matrix)

    # lower bound: sum of node sizes
    if isinstance(node_size, (int, float)):
        sizes = node_size * np.ones((total_nodes))
    elif isinstance(node_size, dict):
        sizes = np.array([node_size[node] if node in node_size else 0. for node in unique_nodes])

    sum_of_node_sizes = sizes[np.newaxis, :] + sizes[:, np.newaxis]
    sum_of_node_sizes -= np.diag(np.diag(sum_of_node_sizes)) # squareform requires zeros on diagonal
    lower_bounds = squareform(sum_of_node_sizes)

    def cost_function(positions):
        return 1. / np.sum((pdist(positions.reshape((-1, 2))))**power)

    def constraint_function(positions):
        positions = np.reshape(positions, (-1, 2))
        return pdist(positions)

    initial_positions = _initialise_geometric_node_layout(edges)
    nonlinear_constraint = NonlinearConstraint(constraint_function, lb=lower_bounds, ub=upper_bounds, jac='2-point')
    result = minimize(cost_function, initial_positions.flatten(), method='SLSQP',
                      jac="2-point", constraints=[nonlinear_constraint], options=dict(maxiter=maximum_iterations))

    if not result.success:
        print("Warning: could not compute valid node positions for the given edge lengths.")
        print(f"scipy.optimize.minimize: {result.message}.")

    node_positions_as_array = result.x.reshape((-1, 2))
    node_positions = dict(zip(unique_nodes, node_positions_as_array))
    return node_positions


def _initialise_geometric_node_layout(edges):
    sources, targets = zip(*edges)
    total_nodes = len(set(sources + targets))
    return np.random.rand(total_nodes, 2)


if __name__ == '__main__':

    import matplotlib.pyplot as plt

    def plot_graph(edges, node_layout):
        # poor man's graph plotting
        fig, ax = plt.subplots()
        for source, target in edges:
            x1, y1 = node_layout[source]
            x2, y2 = node_layout[target]
            ax.plot([x1, x2], [y1, y2], color='darkgray')
        ax.set_aspect('equal')

    ################################################################################
    # triangle with right angle

    edges = [
        (0, 1),
        (1, 2),
        (2, 0)
    ]

    lengths = {
        (0, 1) : 3,
        (1, 2) : 4,
        (2, 0) : 5,
    }

    pos = get_geometric_node_layout(edges, lengths, node_size=0)

    plot_graph(edges, node_layout=pos)

    plt.show()

    ################################################################################
    # square

    edges = [
        (0, 1),
        (1, 2),
        (2, 3),
        (3, 0),
    ]

    lengths = {
        (0, 1) : 0.5,
        (1, 2) : 0.5,
        (2, 3) : 0.5,
        (3, 0) : 0.5,
    }

    pos = get_geometric_node_layout(edges, lengths, node_size=0)

    plot_graph(edges, node_layout=pos)

    plt.show()

Answer 1

此实现计算所有点对的雅可比矩阵。 np 数组向量化代码可能没有优化，因此我欢迎进一步的 cmets 来改进代码。

推导

雅可比矩阵，根据scipy.optimize：

代码

from itertools import combinations

n_pt = len(initial_positions)
n_ptpair = len(upper_bounds) #total number of pointpairs

idx_pts= np.array(list(combinations(range(n_pt),2)))
idx_pt1= np.array(idx_pts[:,0])
idx_pt2= np.array(idx_pts[:,1]) #[[pair_id,x_00,x_01],[1,p_10,p_11],...]

row_idx = np.repeat(np.arange(n_ptpair).reshape(-1,1),2,axis=1)
col1_idx = np.vstack((idx_pt1*2,idx_pt1*2+1)).T
col2_idx = np.vstack((idx_pt2*2,idx_pt2*2+1)).T

def delta_constraint(positions):
  positions = np.reshape(positions, (-1, 2))
  pairdist = constraint_function(positions) #pairwise R2 distance between each point pair with edge
  jac = np.zeros((n_ptpair,2*n_pt)) #(N,(x0,y0,x1,y1,...,xc,yc...,xN,yN))
  jac[row_idx,col1_idx] = (positions[idx_pt1]-positions[idx_pt2])/pairdist.reshape((-1,1))
  jac[row_idx,col2_idx] = -jac[row_idx,col1_idx]
  return jac