-
Notifications
You must be signed in to change notification settings - Fork 29
/
laplacian_funcs.py
131 lines (112 loc) · 4.29 KB
/
laplacian_funcs.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
"""Functions related to getting the laplacian and the right number of pixels after pooling/unpooling.
"""
import numpy as np
import torch
from pygsp.graphs.nngraphs.spherehealpix import SphereHealpix
from pygsp.graphs.nngraphs.sphereicosahedron import SphereIcosahedron
from pygsp.graphs.sphereequiangular import SphereEquiangular
from scipy import sparse
from scipy.sparse import coo_matrix
from deepsphere.utils.samplings import (
equiangular_bandwidth,
equiangular_dimension_unpack,
healpix_resolution_calculator,
icosahedron_nodes_calculator,
icosahedron_order_calculator,
)
def scipy_csr_to_sparse_tensor(csr_mat):
"""Convert scipy csr to sparse pytorch tensor.
Args:
csr_mat (csr_matrix): The sparse scipy matrix.
Returns:
sparse_tensor :obj:`torch.sparse.FloatTensor`: The sparse torch matrix.
"""
coo = coo_matrix(csr_mat)
values = coo.data
indices = np.vstack((coo.row, coo.col))
idx = torch.LongTensor(indices)
vals = torch.FloatTensor(values)
shape = coo.shape
sparse_tensor = torch.sparse.FloatTensor(idx, vals, torch.Size(shape))
sparse_tensor = sparse_tensor.coalesce()
return sparse_tensor
def prepare_laplacian(laplacian):
"""Prepare a graph Laplacian to be fed to a graph convolutional layer.
"""
def estimate_lmax(laplacian, tol=5e-3):
"""Estimate the largest eigenvalue of an operator.
"""
lmax = sparse.linalg.eigsh(laplacian, k=1, tol=tol, ncv=min(laplacian.shape[0], 10), return_eigenvectors=False)
lmax = lmax[0]
lmax *= 1 + 2 * tol # Be robust to errors.
return lmax
def scale_operator(L, lmax, scale=1):
"""Scale the eigenvalues from [0, lmax] to [-scale, scale].
"""
I = sparse.identity(L.shape[0], format=L.format, dtype=L.dtype)
L *= 2 * scale / lmax
L -= I
return L
lmax = estimate_lmax(laplacian)
laplacian = scale_operator(laplacian, lmax)
laplacian = scipy_csr_to_sparse_tensor(laplacian)
return laplacian
def get_icosahedron_laplacians(nodes, depth, laplacian_type):
"""Get the icosahedron laplacian list for a certain depth.
Args:
nodes (int): initial number of nodes.
depth (int): the depth of the UNet.
laplacian_type ["combinatorial", "normalized"]: the type of the laplacian.
Returns:
laps (list): increasing list of laplacians.
"""
laps = []
order = icosahedron_order_calculator(nodes)
for _ in range(depth):
nodes = icosahedron_nodes_calculator(order)
order_initial = icosahedron_order_calculator(nodes)
G = SphereIcosahedron(level=int(order_initial))
G.compute_laplacian(laplacian_type)
laplacian = prepare_laplacian(G.L)
laps.append(laplacian)
order -= 1
return laps[::-1]
def get_healpix_laplacians(nodes, depth, laplacian_type):
"""Get the healpix laplacian list for a certain depth.
Args:
nodes (int): initial number of nodes.
depth (int): the depth of the UNet.
laplacian_type ["combinatorial", "normalized"]: the type of the laplacian.
Returns:
laps (list): increasing list of laplacians.
"""
laps = []
for i in range(depth):
pixel_num = nodes
subdivisions = int(healpix_resolution_calculator(pixel_num)/2**i)
G = SphereHealpix(subdivisions)
G.compute_laplacian(laplacian_type)
laplacian = prepare_laplacian(G.L)
laps.append(laplacian)
return laps[::-1]
def get_equiangular_laplacians(nodes, depth, ratio, laplacian_type):
"""Get the equiangular laplacian list for a certain depth.
Args:
nodes (int): initial number of nodes.
depth (int): the depth of the UNet.
laplacian_type ["combinatorial", "normalized"]: the type of the laplacian.
Returns:
laps (list): increasing list of laplacians
"""
laps = []
pixel_num = nodes
for _ in range(depth):
dim1, dim2 = equiangular_dimension_unpack(pixel_num, ratio)
bw1 = equiangular_bandwidth(dim1)
bw2 = equiangular_bandwidth(dim2)
bw = [bw1, bw2]
G = SphereEquiangular(bandwidth=bw, sampling="SOFT")
G.compute_laplacian(laplacian_type)
laplacian = prepare_laplacian(G.L)
laps.append(laplacian)
return laps[::-1]