# Source code for pygsp.graphs.torus

# -*- coding: utf-8 -*-

from . import Graph

import numpy as np
from scipy import sparse

[docs]class Torus(Graph):
r"""
Create a Torus graph.

Parameters
----------
Nv : int
Number of vertices along the first dimension (default is 16)
Mv : int
Number of vertices along the second dimension (default is Nv)

Examples
--------
>>> from pygsp import graphs
>>> Nv = 32
>>> G = graphs.Torus(Nv=Nv)

References
----------
See :cite:strang1999discrete for more informations.

"""

def __init__(self, Nv=16, Mv=None, **kwargs):

if not Mv:
Mv = Nv

K = 2 * Nv
J = 2 * Mv
i_inds = np.zeros((K*Mv + J*Nv), dtype=float)
j_inds = np.zeros((K*Mv + J*Nv), dtype=float)

tmpK = np.arange(K, dtype=int)
tmpNv1 = np.arange(Nv - 1)
tmpNv = np.arange(Nv)

for i in range(Mv):
i_inds[i*K + tmpK] = i*Nv + \
np.concatenate((np.array([Nv - 1]), tmpNv1, tmpNv))

j_inds[i*K + tmpK] = i*Nv + \
np.concatenate((tmpNv, np.array([Nv - 1]), tmpNv1))

tmp2Nv = np.arange(2*Nv, dtype=int)

for i in range(Mv - 1):
i_inds[K*Mv + i*2*Nv + tmp2Nv] = \
np.concatenate((i*Nv + tmpNv, (i + 1)*Nv + tmpNv))

j_inds[K*Mv + i*2*Nv + tmp2Nv] = \
np.concatenate(((i + 1)*Nv + tmpNv, i*Nv + tmpNv))

i_inds[K*Mv + (Mv - 1)*2*Nv + tmp2Nv] = \
np.concatenate((tmpNv, (Mv - 1)*Nv + tmpNv))

j_inds[K*Mv + (Mv - 1)*2*Nv + tmp2Nv] = \
np.concatenate(((Mv - 1)*Nv + tmpNv, tmpNv))

W = sparse.csc_matrix((np.ones((K*Mv + J*Nv)), (i_inds, j_inds)),
shape=(Mv*Nv, Mv*Nv))

# Create coordinate
T = 1.5 + np.sin(np.arange(Mv)*2*np.pi/Mv).reshape(1, Mv)
U = np.cos(np.arange(Mv)*2*np.pi/Mv).reshape(1, Mv)
xtmp = np.cos(np.arange(Nv).reshape(Nv, 1)*2*np.pi/Nv)*T
ytmp = np.sin(np.arange(Nv).reshape(Nv, 1)*2*np.pi/Nv)*T
ztmp = np.kron(np.ones((Nv, 1)), U)
coords = np.concatenate((np.reshape(xtmp, (Mv*Nv, 1), order='F'),
np.reshape(ytmp, (Mv*Nv, 1), order='F'),
np.reshape(ztmp, (Mv*Nv, 1), order='F')),
axis=1)
self.Nv = Nv
self.Mv = Nv

plotting = {"vertex_size": 30,
"limits": np.array([-2.5, 2.5, -2.5, 2.5, -2.5, 2.5])}

super(Torus, self).__init__(W=W, gtype='Torus', coords=coords,
plotting=plotting, **kwargs)