Min Cost Flow

def calc(self, maxiter=100, fixedprec=1e9):
        """Min Cost Flow"""
        source_data_holder = []
        N = self.targets.shape[0]
        K = self.origins.shape[0]

        # dict of labels for each target node
        M, demand = self._get_demand_graph()

        max_dist_trip = 400  # kilometers
        cost_holder = []
        itercnt = 0

        while True:
            itercnt += 1
            logging.info(f'Iter count: {itercnt}')

            # Setup the graph
            g = nx.DiGraph()

            self.targets = self.targets.sort_values('labels').reset_index(drop=True)
            # Supply of 1 (i.e. demand = -1) means that it can only be connected to one node
            g.add_nodes_from(self.targets['target_id'], demand=-1) # points

            for idx in self.nearest_targets.origin_id:
                g.add_node(int(idx), demand=demand[idx])

            ### Create the cluster centers calculate a distance cost
            cost_dist = dist_vect(
                np.tile(self.targets['lng'].values, K),
                np.tile(self.targets['lat'].values, K),
                np.repeat(self.origins['lng'].values, N),
                np.repeat(self.origins['lat'].values, N)
            )[:, np.newaxis].T

            scaler_dist = MinMaxScaler()
            cost_dist_trans = scaler_dist.fit_transform(cost_dist.T).T

            # Penalty for distances too large
            cost_dist_trans[cost_dist > max_dist_trip] = 10

            # Create the in-cluster sales and calculate the total volume of sales generated
            # TODO: rename this to something more generic, like cluster_demanded
            cluster_sales = self.targets.groupby('labels').sum()[self.demand_col][:, np.newaxis]
            D = cluster_sales.shape[1]

            cost_sales = abs(
                np.array([
                    np.linalg.norm(
                            np.repeat(cluster_sales, N)[:, np.newaxis] \
                            - np.tile(cluster_sales.mean(), (K * N))[:,np.newaxis],
                            axis=1
                        )
                    ])
                )

            scaler_sales = MinMaxScaler()
            cost_sales = scaler_sales.fit_transform(cost_sales.T).T

            # Total cost TO CHANGE??
            cost_total = cost_dist_trans + cost_sales

            cost_holder.append(sum(cost_total[0]))


            # Create the edges of points to centers
            data_to_center_edges = np.concatenate(
                (
                    np.tile(self.targets['target_id'], K).T[:, np.newaxis],
                    np.array([np.tile(int(i+1), self.targets.shape[0]) for i in range(K)]).reshape(self.targets.shape[0] * K, 1),
                    cost_total.T * 1e5
                ),
                axis=1
            ).astype(np.uint64)

            # Add these edges to the graph
            g.add_weighted_edges_from(data_to_center_edges)

            # Add the extra balance node
            # To balance out the network, we create an extra node that is:
            # -(K*(-1)+sum(demand_per_node))
            a = 99999
            g.add_node(a, demand=self.targets.shape[0] - np.sum(list(demand.values())))

            C_to_a_edges = np.concatenate(
                (
                    np.array([int(i + 1) for i in range(K)]).T[:, np.newaxis],
                    np.tile([[a, ]], K).T
                ),
                axis=1
            )
            g.add_edges_from(C_to_a_edges)

            # Calculate the minimum flow cost
            f = nx.min_cost_flow(g)

            # Update the new labels
            M_new = {}
            p = {}

            for i in list(g.nodes)[:-1]:
                # Sorts all the items in the dictionary and picks the cluster
                # with label = 1
                p = sorted(f[i].items(), key=lambda x: x[1])[-1][0]
                M_new[i] = p

            # Update the new labels in the df
            self.targets['labels'] = self.targets.apply(lambda x: M_new[x['target_id']], axis=1)

            # Set the capacity for all edges
            # TO DO: Figure how/whether we need to properly set a capacity for the edges.
            C = 50
            nx.set_edge_attributes(g, C, 'capacity')

            # Test whether we can stop
            # stop condition
            if np.all(M_new == M):
                print("All same")
                self.results = {
                    'dict_graph': M,
                    'min_cost_flow': f,
                    'nxgraph': g,
                    'model_labels': self.targets,
                    'costs': cost_holder
                }
                return True

            M = M_new

            source_data_holder.append(self.targets['labels'].values)
            if maxiter is not None and itercnt >= maxiter:
                # Max iterations reached
                self.results = {
                    'dict_graph': M,
                    'min_cost_flow': f,
                    'nxgraph': g,
                    'model_labels': self.targets,
                    'costs': cost_holder
                }
                return True