gvg.py

import numpy as np
import time
from tqdm import tqdm
import scipy.sparse as sp
from scipy.sparse.csgraph import dijkstra
import math
from scipy.stats import multivariate_normal
import multiprocessing as mp


class GVG:
    def __init__(self, W, R, scalar=True, logging=True, mode="exists"): # mode exists or all
        self.W = W
        self.R = R
        self.mode = mode
        self.logging = logging
        self.tol = 1e-9
        self.f_vis = self.check_vis_scalar if scalar else self.check_vis_vectoral
        self.unweighted = self.is_unweighted()
        self.symmetric = self.is_symmetric()
        
        self.tau = 0.5
        if self.mode == "fraction" and (self.tau <= 0 or self.tau > 1):
            raise ValueError("For fractional visibility, tau must satisfy 0 < tau <= 1.")

        self.build_neighbors()

    def is_symmetric(self):
        if sp.issparse(self.W):
            return (self.W != self.W.T).nnz == 0
        else:
            return np.all(self.W == self.W.T)


    def is_unweighted(self):
        if sp.issparse(self.W):
            M = self.W.data
            return np.all((M == 1) | (M == 0))
        else:
            return np.all((self.W == 1) | (self.W == 0))


    def build_neighbors(self):
        if sp.issparse(self.W):
            W = self.W.tocsr()
            self.neigh = []
            self.wts = []
            for u in range(W.shape[0]):
                cols = W.indices[W.indptr[u]:W.indptr[u+1]]
                vals = W.data[W.indptr[u]:W.indptr[u+1]]
                self.neigh.append(cols)
                self.wts.append(vals)
        else:
            W = np.asarray(self.W)
            self.neigh = []
            self.wts = []
            for u in range(W.shape[0]):
                cols = np.where(W[u] != 0)[0]
                self.neigh.append(cols)
                self.wts.append(W[u, cols])


    def comp_dijkstra(self):
        dist = dijkstra(
            self.W,
            unweighted=self.unweighted,
            return_predecessors=False,
            limit=self.R
        )
        return dist


    def predec(self, dist_row):
        n = len(dist_row)
        preds = [[] for _ in range(n)]
        d = dist_row
    
        for u in range(n):
            du = d[u]
            if not np.isfinite(du):
                continue
    
            for idx, v in enumerate(self.neigh[u]):
                dv = d[v]
                if not np.isfinite(dv):
                    continue
    
                w = 1.0 if self.unweighted else self.wts[u][idx]
                if abs((du + w) - dv) <= self.tol:
                    preds[v].append(u)
        return preds


    def check_vis_scalar(self, x, path, d_path):
        # d_path is accumulated distance along the path
        # x is vector in R^n
        if len(path) <= 2:
            return True
        X_s = x[path[0]]
        X_t = x[path[-1]]
        d_st = d_path[-1] - d_path[0]   
         
        for u in range(1, len(path) - 1):
            X_u = x[path[u]]
            d_su = d_path[u] - d_path[0] # dist up to node u from start s
            a_su = d_su / d_st
            L = (1 - a_su) * X_s + a_su * X_t
            if X_u >= L:
                return False
        return True


    def check_vis_vectoral(self, x, path, d_path, eps=1e-12):

        if len(path) <= 2:
            return True
    
        s = path[0]
        t = path[-1]
    
        X_s = x[s]
        X_t = x[t]
    
        d_st = d_path[-1] - d_path[0]
        if d_st <= 0:
            return True  
    
        X_ss = np.linalg.norm(X_s)
        
        if X_ss <= eps:
            return False
    
        P_s = X_ss
        P_t = np.dot(X_s, X_t) / X_ss
    
        for u in range(1, len(path) - 1):
            node_u = path[u]
            X_u = x[node_u]
            P_u = np.dot(X_s, X_u) / X_ss
    
            d_su = d_path[u] - d_path[0]
            a_su = d_su / d_st
    
            L = (1.0 - a_su) * P_s + a_su * P_t
    
            if P_u >= L:
                return False
    
        return True


    def soft_visibility(self, s, t, preds, x, dist_row):
        # return True immediately if we find ONE shortest path that is visible
        path_rev = [t]
        
        def dfs(v):
            if v == s:
                path = list(reversed(path_rev))
                d_path = [dist_row[w] for w in path]
                return self.f_vis(x, path, d_path)
                
            if len(preds[v]) == 0:
                return False
            
            for u in preds[v]:
                path_rev.append(u)
                ok = dfs(u)
                path_rev.pop()
                if ok:
                    return True
            return False

        return dfs(t)


    def strict_visibility(self, s, t, preds, x, dist_row):
        # return False immediately if we find ANY shortest path that is NOT visible
        path_rev = [t]

        def dfs(v):
            if v == s:
                path = list(reversed(path_rev))
                d_path = [dist_row[w] for w in path]
                return self.f_vis(x, path, d_path)
            
            if len(preds[v]) == 0:
                return False
            
            for u in preds[v]:
                path_rev.append(u)
                ok = dfs(u)
                path_rev.pop()
                if not ok:
                    return False
            return True

        return dfs(t)


    def frac_visibility(self, s, t, preds, x, dist_row):
        # return True immediately if we accumulate at least ceil[tau * M_st] number of paths that are visible
        M_st, _ = self.path_cnt(s, t, preds)
        if M_st <= 0:
            return False
    
        thr = int(math.ceil(self.tau * M_st))
        path_rev = [t]
        acc_vis = 0
    
        def dfs(v):
            nonlocal acc_vis
            if v == s:
                path = list(reversed(path_rev))
                d_path = [dist_row[w] for w in path]
                ok = self.f_vis(x, path, d_path)
                if ok:
                    acc_vis += 1
                    if acc_vis >= thr:
                        return True
                return False
    
            if len(preds[v]) == 0:
                return False
            
            for u in preds[v]:
                path_rev.append(u)
                hit = dfs(u)
                path_rev.pop()
                if hit:
                    return True
            return False
    
        return dfs(t)


    def path_cnt(self, s, t, preds, cap=10**8):
        n = len(preds)
        cnt = -np.ones(n, dtype=np.int64)
        hop = -np.ones(n, dtype=np.int64)
    
        def dfs(v):
            if v == s:
                return 1, 0
            if cnt[v] >= 0:
                return cnt[v], hop[v]
    
            if len(preds[v]) == 0:
                c, h = 0, 10**8
            else:
                total = 0
                best_h = 10**8
                for u in preds[v]:
                    cu, hu = dfs(u)
                    total += cu
                    if total >= cap:
                        total = cap
                    if hu + 1 < best_h:
                        best_h = hu + 1
                c, h = total, best_h
    
            cnt[v] = c
            hop[v] = h
            return c, h
    
        return dfs(t)

        
    def build(self, x):
        start_time = time.time()
        N = len(x)
        G_v = np.zeros((N, N), dtype=np.int8)

        dist = self.comp_dijkstra()

        for s in tqdm(range(N), desc="Visibility Computation", disable=not self.logging):
            dist_row = dist[s, :]
            preds = self.predec(dist_row)
            targets = np.where(np.isfinite(dist_row) & (dist_row > 0) & (dist_row <= self.R))[0]
            
            for t in targets:
                if len(preds[t]) == 0:
                    continue

                if self.mode == "exists":
                    ok = self.soft_visibility(s, t, preds, x, dist_row)
                elif self.mode == "all":
                    ok = self.strict_visibility(s, t, preds, x, dist_row)
                elif self.mode == "fraction":
                    ok = self.frac_visibility(s, t, preds, x, dist_row)
                else:
                    raise ValueError("mode must be 'exists' or 'all'")
                G_v[s, t] = 1 if ok else 0

        if self.logging:
            print("Unweighted:", self.unweighted, "| symmetric:", self.symmetric, "| mode:", self.mode)
            print(f"Total time: {time.time() - start_time:.2f}s")
            print("G_v is symmetric?:", np.all(G_v == G_v.T))
        return G_v, dist


    def build_parallel(self, x, n_jobs=16, chunk_size=50, return_dense=True):
        start_time = time.time()
        N = len(x)
    
        dist = self.comp_dijkstra()  
        sources = np.arange(N, dtype=np.int32)
        chunks = [sources[i:i+chunk_size] for i in range(0, N, chunk_size)]
    
        try:
            ctx = mp.get_context("fork")
        except ValueError:
            ctx = mp.get_context()  
    
        n_jobs = int(n_jobs)
        if n_jobs <= 0:
            n_jobs = 1
    
        with ctx.Pool(
            processes=n_jobs,
            initializer=_init_worker,
            initargs=(self, dist, x, N)
        ) as pool:

            results = pool.map(_process_sources, chunks)
    
        all_rows = []
        all_cols = []
        for r, c in results:
            if r:
                all_rows.append(np.asarray(r, dtype=np.int32))
                all_cols.append(np.asarray(c, dtype=np.int32))
    
        if len(all_rows) == 0:
            G_sparse = sp.csr_matrix((N, N), dtype=np.int8)
        else:
            rr = np.concatenate(all_rows)
            cc = np.concatenate(all_cols)
            data = np.ones(rr.shape[0], dtype=np.int8)
            G_sparse = sp.csr_matrix((data, (rr, cc)), shape=(N, N), dtype=np.int8)
    
        if self.logging:
            print(f"[build_parallel] jobs={n_jobs} chunk={chunk_size} time={time.time()-start_time:.2f}s")
            diff = (G_sparse != G_sparse.T).nnz
            print("[build_parallel] G_v symmetric?:", diff == 0)

        if return_dense:
            G_dense = G_sparse.toarray().astype(np.int8)
            return G_dense, dist
        else:
            return G_sparse, dist


_GVG_obj = None
_DIST = None
_X = None
_N = None

def _init_worker(gvg_obj, dist, x, n):
    global _GVG_obj, _DIST, _X, _N
    _GVG_obj = gvg_obj
    _DIST = dist
    _X = x
    _N = n

def _process_sources(s_list):
    gvg = _GVG_obj
    dist = _DIST
    x = _X

    rows = []
    cols = []

    for s in s_list:
        dist_row = dist[s, :]
        preds = gvg.predec(dist_row)
        targets = np.where(np.isfinite(dist_row) & (dist_row > 0) & (dist_row <= gvg.R))[0]

        for t in targets:
            if len(preds[t]) == 0:
                continue

            if gvg.mode == "exists":
                ok = gvg.soft_visibility(s, t, preds, x, dist_row)
            elif gvg.mode == "all":
                ok = gvg.strict_visibility(s, t, preds, x, dist_row)
            elif gvg.mode == "fraction":
                ok = gvg.frac_visibility(s, t, preds, x, dist_row)
            else:
                raise ValueError("mode must be 'exists', 'all', or 'fraction'")

            if ok:
                rows.append(s)
                cols.append(t)

    return rows, cols