Upload 30 files

cortex_DIM/configs/convnets.py

@@ -0,0 +1,98 @@
+'''Basic convnet hyperparameters.
+
+conv_args are in format (dim_h, f_size, stride, pad batch_norm, dropout, nonlinearity, pool)
+fc_args are in format (dim_h, batch_norm, dropout, nonlinearity)
+
+'''
+
+from cortex_DIM.nn_modules.encoder import ConvnetEncoder, FoldedConvnetEncoder
+
+
+# Basic DCGAN-like encoders
+
+_basic28x28 = dict(
+    Encoder=ConvnetEncoder,
+    conv_args=[(64, 5, 2, 2, True, False, 'ReLU', None),
+               (128, 5, 2, 2, True, False, 'ReLU', None)],
+    fc_args=[(1024, True, False, 'ReLU', None)],
+    local_idx=1,
+    fc_idx=0
+)
+
+_basic32x32 = dict(
+    Encoder=ConvnetEncoder,
+    conv_args=[(64, 4, 2, 1, True, False, 'ReLU', None),
+               (128, 4, 2, 1, True, False, 'ReLU', None),
+               (256, 4, 2, 1, True, False, 'ReLU', None)],
+    fc_args=[(1024, True, False, 'ReLU')],
+    local_idx=1,
+    conv_idx=2,
+    fc_idx=0
+)
+
+_basic64x64 = dict(
+    Encoder=ConvnetEncoder,
+    conv_args=[(64, 4, 2, 1, True, False, 'ReLU', None),
+               (128, 4, 2, 1, True, False, 'ReLU', None),
+               (256, 4, 2, 1, True, False, 'ReLU', None),
+               (512, 4, 2, 1, True, False, 'ReLU', None)],
+    fc_args=[(1024, True, False, 'ReLU')],
+    local_idx=2,
+    conv_idx=3,
+    fc_idx=0
+)
+
+# Alexnet-like encoders
+
+_alex64x64 = dict(
+    Encoder=ConvnetEncoder,
+    conv_args=[(96, 3, 1, 1, True, False, 'ReLU', ('MaxPool2d', 3, 2)),
+               (192, 3, 1, 1, True, False, 'ReLU', ('MaxPool2d', 3, 2)),
+               (384, 3, 1, 1, True, False, 'ReLU', None),
+               (384, 3, 1, 1, True, False, 'ReLU', None),
+               (192, 3, 1, 1, True, False, 'ReLU', ('MaxPool2d', 3, 2))],
+    fc_args=[(4096, True, False, 'ReLU'),
+             (4096, True, False, 'ReLU')],
+    local_idx=2,
+    conv_idx=4,
+    fc_idx=1
+)
+
+_foldalex64x64 = dict(
+    Encoder=FoldedConvnetEncoder,
+    crop_size=16,
+    conv_args=[(96, 3, 1, 1, True, False, 'ReLU', ('MaxPool2d', 3, 2)),
+               (192, 3, 1, 1, True, False, 'ReLU', ('MaxPool2d', 3, 2)),
+               (384, 3, 1, 1, True, False, 'ReLU', None),
+               (384, 3, 1, 1, True, False, 'ReLU', None),
+               (192, 3, 1, 1, True, False, 'ReLU', ('MaxPool2d', 3, 2))],
+    fc_args=[(4096, True, False, 'ReLU'),
+             (4096, True, False, 'ReLU')],
+    local_idx=4,
+    fc_idx=1
+)
+
+_foldmultialex64x64 = dict(
+    Encoder=FoldedConvnetEncoder,
+    crop_size=16,
+    conv_args=[(96, 3, 1, 1, True, False, 'ReLU', ('MaxPool2d', 3, 2)),
+               (192, 3, 1, 1, True, False, 'ReLU', ('MaxPool2d', 3, 2)),
+               (384, 3, 1, 1, True, False, 'ReLU', None),
+               (384, 3, 1, 1, True, False, 'ReLU', None),
+               (192, 3, 1, 1, True, False, 'ReLU', ('MaxPool2d', 3, 2)),
+               (192, 3, 1, 0, True, False, 'ReLU', None),
+               (192, 1, 1, 0, True, False, 'ReLU', None)],
+    fc_args=[(4096, True, False, 'ReLU')],
+    local_idx=4,
+    multi_idx=6,
+    fc_idx=1
+)
+
+configs = dict(
+    basic28x28=_basic28x28,
+    basic32x32=_basic32x32,
+    basic64x64=_basic64x64,
+    alex64x64=_alex64x64,
+    foldalex64x64=_foldalex64x64,
+    foldmultialex64x64=_foldmultialex64x64
+)

cortex_DIM/configs/resnets.py

@@ -0,0 +1,151 @@
+"""Configurations for ResNets
+
+"""
+
+from cortex_DIM.nn_modules.encoder import ResnetEncoder, FoldedResnetEncoder
+
+
+_resnet19_32x32 = dict(
+    Encoder=ResnetEncoder,
+    conv_before_args=[(64, 3, 2, 1, True, False, 'ReLU', None)],
+    res_args=[
+        ([(64, 1, 1, 0, True, False, 'ReLU', None),
+          (64, 3, 1, 1, True, False, 'ReLU', None),
+          (64 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(64, 1, 1, 0, True, False, 'ReLU', None),
+          (64, 3, 1, 1, True, False, 'ReLU', None),
+          (64 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(128, 1, 1, 0, True, False, 'ReLU', None),
+          (128, 3, 2, 1, True, False, 'ReLU', None),
+          (128 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(128, 1, 1, 0, True, False, 'ReLU', None),
+          (128, 3, 1, 1, True, False, 'ReLU', None),
+          (128 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(256, 1, 1, 0, True, False, 'ReLU', None),
+          (256, 3, 2, 1, True, False, 'ReLU', None),
+          (256 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(256, 1, 1, 0, True, False, 'ReLU', None),
+          (256, 3, 1, 1, True, False, 'ReLU', None),
+          (256 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1)
+    ],
+    fc_args=[(1024, True, False, 'ReLU')],
+    local_idx=4,
+    fc_idx=0
+)
+
+_foldresnet19_32x32 = dict(
+    Encoder=FoldedResnetEncoder,
+    crop_size=8,
+    conv_before_args=[(64, 3, 2, 1, True, False, 'ReLU', None)],
+    res_args=[
+        ([(64, 1, 1, 0, True, False, 'ReLU', None),
+          (64, 3, 1, 1, True, False, 'ReLU', None),
+          (64 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(64, 1, 1, 0, True, False, 'ReLU', None),
+          (64, 3, 1, 1, True, False, 'ReLU', None),
+          (64 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(128, 1, 1, 0, True, False, 'ReLU', None),
+          (128, 3, 2, 1, True, False, 'ReLU', None),
+          (128 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(128, 1, 1, 0, True, False, 'ReLU', None),
+          (128, 3, 1, 1, True, False, 'ReLU', None),
+          (128 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(256, 1, 1, 0, True, False, 'ReLU', None),
+          (256, 3, 2, 1, True, False, 'ReLU', None),
+          (256 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(256, 1, 1, 0, True, False, 'ReLU', None),
+          (256, 3, 1, 1, True, False, 'ReLU', None),
+          (256 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1)
+    ],
+    fc_args=[(1024, True, False, 'ReLU')],
+    local_idx=6,
+    fc_idx=0
+)
+
+_resnet34_32x32 = dict(
+    Encoder=ResnetEncoder,
+    conv_before_args=[(64, 3, 2, 1, True, False, 'ReLU', None)],
+    res_args=[
+        ([(64, 1, 1, 0, True, False, 'ReLU', None),
+          (64, 3, 1, 1, True, False, 'ReLU', None),
+          (64 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(64, 1, 1, 0, True, False, 'ReLU', None),
+          (64, 3, 1, 1, True, False, 'ReLU', None),
+          (64 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         2),
+        ([(128, 1, 1, 0, True, False, 'ReLU', None),
+          (128, 3, 2, 1, True, False, 'ReLU', None),
+          (128 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(128, 1, 1, 0, True, False, 'ReLU', None),
+          (128, 3, 1, 1, True, False, 'ReLU', None),
+          (128 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         5),
+        ([(256, 1, 1, 0, True, False, 'ReLU', None),
+          (256, 3, 2, 1, True, False, 'ReLU', None),
+          (256 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(256, 1, 1, 0, True, False, 'ReLU', None),
+          (256, 3, 1, 1, True, False, 'ReLU', None),
+          (256 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         2)
+    ],
+    fc_args=[(1024, True, False, 'ReLU')],
+    local_idx=2,
+    fc_idx=0
+)
+
+_foldresnet34_32x32 = dict(
+    Encoder=FoldedResnetEncoder,
+    crop_size=8,
+    conv_before_args=[(64, 3, 2, 1, True, False, 'ReLU', None)],
+    res_args=[
+        ([(64, 1, 1, 0, True, False, 'ReLU', None),
+          (64, 3, 1, 1, True, False, 'ReLU', None),
+          (64 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(64, 1, 1, 0, True, False, 'ReLU', None),
+          (64, 3, 1, 1, True, False, 'ReLU', None),
+          (64 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         2),
+        ([(128, 1, 1, 0, True, False, 'ReLU', None),
+          (128, 3, 2, 1, True, False, 'ReLU', None),
+          (128 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(128, 1, 1, 0, True, False, 'ReLU', None),
+          (128, 3, 1, 1, True, False, 'ReLU', None),
+          (128 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         5),
+        ([(256, 1, 1, 0, True, False, 'ReLU', None),
+          (256, 3, 2, 1, True, False, 'ReLU', None),
+          (256 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         1),
+        ([(256, 1, 1, 0, True, False, 'ReLU', None),
+          (256, 3, 1, 1, True, False, 'ReLU', None),
+          (256 * 4, 1, 1, 0, True, False, 'ReLU', None)],
+         2)
+    ],
+    fc_args=[(1024, True, False, 'ReLU')],
+    local_idx=12,
+    fc_idx=0
+)
+
+configs = dict(
+    resnet19_32x32=_resnet19_32x32,
+    resnet34_32x32=_resnet34_32x32,
+    foldresnet19_32x32=_foldresnet19_32x32,
+    foldresnet34_32x32=_foldresnet34_32x32
+)

cortex_DIM/functions/dim_losses.py

@@ -0,0 +1,224 @@
+'''cortex_DIM losses.
+
+'''
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+from cortex_DIM.functions.gan_losses import get_positive_expectation, get_negative_expectation
+
+
+def fenchel_dual_loss(l, g, measure=None):
+    '''Computes the f-divergence distance between positive and negative joint distributions.
+
+    Divergences supported are Jensen-Shannon `JSD`, `GAN` (equivalent to JSD),
+    Squared Hellinger `H2`, Chi-squeared `X2`, `KL`, and reverse KL `RKL`.
+
+    Args:
+        l: Local feature map.
+        g: Global features.
+        measure: f-divergence measure.
+
+    Returns:
+        torch.Tensor: Loss.
+
+    '''
+    N, local_units, n_locs = l.size()
+    l = l.permute(0, 2, 1)
+    l = l.reshape(-1, local_units)
+
+    u = torch.mm(g, l.t())
+    u = u.reshape(N, N, -1)
+    mask = torch.eye(N).cuda()
+    n_mask = 1 - mask
+
+    E_pos = get_positive_expectation(u, measure, average=False).mean(2)
+    E_neg = get_negative_expectation(u, measure, average=False).mean(2)
+    E_pos = (E_pos * mask).sum() / mask.sum()
+    E_neg = (E_neg * n_mask).sum() / n_mask.sum()
+    loss = E_neg - E_pos
+    return loss
+
+
+def multi_fenchel_dual_loss(l, m, measure=None):
+    '''Computes the f-divergence distance between positive and negative joint distributions.
+
+    Used for multiple globals.
+
+    Divergences supported are Jensen-Shannon `JSD`, `GAN` (equivalent to JSD),
+    Squared Hellinger `H2`, Chi-squeared `X2`, `KL`, and reverse KL `RKL`.
+
+    Args:
+        l: Local feature map.
+        m: Multiple globals feature map.
+        measure: f-divergence measure.
+
+    Returns:
+        torch.Tensor: Loss.
+
+    '''
+    N, units, n_locals = l.size()
+    n_multis = m.size(2)
+
+    l = l.view(N, units, n_locals)
+    l = l.permute(0, 2, 1)
+    l = l.reshape(-1, units)
+
+    m = m.view(N, units, n_multis)
+    m = m.permute(0, 2, 1)
+    m = m.reshape(-1, units)
+
+    u = torch.mm(m, l.t())
+    u = u.reshape(N, n_multis, N, n_locals).permute(0, 2, 3, 1)
+
+    mask = torch.eye(N).cuda()
+    n_mask = 1 - mask
+
+    E_pos = get_positive_expectation(u, measure, average=False).mean(2).mean(2)
+    E_neg = get_negative_expectation(u, measure, average=False).mean(2).mean(2)
+    E_pos = (E_pos * mask).sum() / mask.sum()
+    E_neg = (E_neg * n_mask).sum() / n_mask.sum()
+    loss = E_neg - E_pos
+    return loss
+
+
+def nce_loss(l, g):
+    '''Computes the noise contrastive estimation-based loss.
+
+    Args:
+        l: Local feature map.
+        g: Global features.
+
+    Returns:
+        torch.Tensor: Loss.
+
+    '''
+    N, local_units, n_locs = l.size()
+    l_p = l.permute(0, 2, 1)
+    u_p = torch.matmul(l_p, g.unsqueeze(dim=2))
+
+    l_n = l_p.reshape(-1, local_units)
+    u_n = torch.mm(g, l_n.t())
+    u_n = u_n.reshape(N, N, n_locs)
+
+    mask = torch.eye(N).unsqueeze(dim=2).cuda()
+    n_mask = 1 - mask
+
+    u_n = (n_mask * u_n) - (10. * (1 - n_mask))  # mask out "self" examples
+    u_n = u_n.reshape(N, -1).unsqueeze(dim=1).expand(-1, n_locs, -1)
+
+    pred_lgt = torch.cat([u_p, u_n], dim=2)
+    pred_log = F.log_softmax(pred_lgt, dim=2)
+    loss = -pred_log[:, :, 0].mean()
+    return loss
+
+
+def multi_nce_loss(l, m):
+    '''
+
+    Used for multiple globals.
+
+    Args:
+        l: Local feature map.
+        m: Multiple globals feature map.
+
+    Returns:
+        torch.Tensor: Loss.
+
+    '''
+    N, units, n_locals = l.size()
+    _, _ , n_multis = m.size()
+
+    l = l.view(N, units, n_locals)
+    m = m.view(N, units, n_multis)
+    l_p = l.permute(0, 2, 1)
+    m_p = m.permute(0, 2, 1)
+    u_p = torch.matmul(l_p, m).unsqueeze(2)
+
+    l_n = l_p.reshape(-1, units)
+    m_n = m_p.reshape(-1, units)
+    u_n = torch.mm(m_n, l_n.t())
+    u_n = u_n.reshape(N, n_multis, N, n_locals).permute(0, 2, 3, 1)
+
+    mask = torch.eye(N)[:, :, None, None].cuda()
+    n_mask = 1 - mask
+
+    u_n = (n_mask * u_n) - (10. * (1 - n_mask))  # mask out "self" examples
+    u_n = u_n.reshape(N, N * n_locals, n_multis).unsqueeze(dim=1).expand(-1, n_locals, -1, -1)
+
+    pred_lgt = torch.cat([u_p, u_n], dim=2)
+    pred_log = F.log_softmax(pred_lgt, dim=2)
+    loss = -pred_log[:, :, 0].mean()
+
+    return loss
+
+
+def donsker_varadhan_loss(l, g):
+    '''
+
+    Args:
+        l: Local feature map.
+        g: Global features.
+
+    Returns:
+        torch.Tensor: Loss.
+
+    '''
+    N, local_units, n_locs = l.size()
+    l = l.permute(0, 2, 1)
+    l = l.reshape(-1, local_units)
+
+    u = torch.mm(g, l.t())
+    u = u.reshape(N, N, n_locs)
+
+    mask = torch.eye(N).cuda()
+    n_mask = (1 - mask)[:, :, None]
+
+    E_pos = (u.mean(2) * mask).sum() / mask.sum()
+
+    u -= 100 * (1 - n_mask)
+    u_max = torch.max(u)
+    E_neg = torch.log((n_mask * torch.exp(u - u_max)).sum() + 1e-6) + u_max - math.log(n_mask.sum())
+    loss = E_neg - E_pos
+    return loss
+
+
+def multi_donsker_varadhan_loss(l, m):
+    '''
+
+    Used for multiple globals.
+
+    Args:
+        l: Local feature map.
+        m: Multiple globals feature map.
+
+    Returns:
+        torch.Tensor: Loss.
+
+    '''
+    N, units, n_locals = l.size()
+    n_multis = m.size(2)
+
+    l = l.view(N, units, n_locals)
+    l = l.permute(0, 2, 1)
+    l = l.reshape(-1, units)
+
+    m = m.view(N, units, n_multis)
+    m = m.permute(0, 2, 1)
+    m = m.reshape(-1, units)
+
+    u = torch.mm(m, l.t())
+    u = u.reshape(N, n_multis, N, n_locals).permute(0, 2, 3, 1)
+
+    mask = torch.eye(N).cuda()
+    n_mask = 1 - mask
+
+    E_pos = (u.mean(2) * mask).sum() / mask.sum()
+
+    u -= 100 * (1 - n_mask)
+    u_max = torch.max(u)
+    E_neg = torch.log((n_mask * torch.exp(u - u_max)).sum() + 1e-6) + u_max - math.log(n_mask.sum())
+    loss = E_neg - E_pos
+    return loss

cortex_DIM/functions/gan_losses.py

@@ -0,0 +1,95 @@
+"""
+
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+from cortex_DIM.functions.misc import log_sum_exp
+
+
+def raise_measure_error(measure):
+    supported_measures = ['GAN', 'JSD', 'X2', 'KL', 'RKL', 'DV', 'H2', 'W1']
+    raise NotImplementedError(
+        'Measure `{}` not supported. Supported: {}'.format(measure,
+                                                           supported_measures))
+
+
+def get_positive_expectation(p_samples, measure, average=True):
+    """Computes the positive part of a divergence / difference.
+
+    Args:
+        p_samples: Positive samples.
+        measure: Measure to compute for.
+        average: Average the result over samples.
+
+    Returns:
+        torch.Tensor
+
+    """
+    log_2 = math.log(2.)
+
+    if measure == 'GAN':
+        Ep = - F.softplus(-p_samples)
+    elif measure == 'JSD':
+        Ep = log_2 - F.softplus(- p_samples)
+    elif measure == 'X2':
+        Ep = p_samples ** 2
+    elif measure == 'KL':
+        Ep = p_samples + 1.
+    elif measure == 'RKL':
+        Ep = -torch.exp(-p_samples)
+    elif measure == 'DV':
+        Ep = p_samples
+    elif measure == 'H2':
+        Ep = 1. - torch.exp(-p_samples)
+    elif measure == 'W1':
+        Ep = p_samples
+    else:
+        raise_measure_error(measure)
+
+    if average:
+        return Ep.mean()
+    else:
+        return Ep
+
+
+def get_negative_expectation(q_samples, measure, average=True):
+    """Computes the negative part of a divergence / difference.
+
+    Args:
+        q_samples: Negative samples.
+        measure: Measure to compute for.
+        average: Average the result over samples.
+
+    Returns:
+        torch.Tensor
+
+    """
+    log_2 = math.log(2.)
+
+    if measure == 'GAN':
+        Eq = F.softplus(-q_samples) + q_samples
+    elif measure == 'JSD':
+        Eq = F.softplus(-q_samples) + q_samples - log_2
+    elif measure == 'X2':
+        Eq = -0.5 * ((torch.sqrt(q_samples ** 2) + 1.) ** 2)
+    elif measure == 'KL':
+        Eq = torch.exp(q_samples)
+    elif measure == 'RKL':
+        Eq = q_samples - 1.
+    elif measure == 'DV':
+        Eq = log_sum_exp(q_samples, 0) - math.log(q_samples.size(0))
+    elif measure == 'H2':
+        Eq = torch.exp(q_samples) - 1.
+    elif measure == 'W1':
+        Eq = q_samples
+    else:
+        raise_measure_error(measure)
+
+    if average:
+        return Eq.mean()
+    else:
+        return Eq

cortex_DIM/functions/misc.py

@@ -0,0 +1,39 @@
+"""Miscilaneous functions.
+
+"""
+
+import torch
+
+
+def log_sum_exp(x, axis=None):
+    """Log sum exp function
+
+    Args:
+        x: Input.
+        axis: Axis over which to perform sum.
+
+    Returns:
+        torch.Tensor: log sum exp
+
+    """
+    x_max = torch.max(x, axis)[0]
+    y = torch.log((torch.exp(x - x_max)).sum(axis)) + x_max
+    return y
+
+
+def random_permute(X):
+    """Randomly permutes a tensor.
+
+    Args:
+        X: Input tensor.
+
+    Returns:
+        torch.Tensor
+
+    """
+    X = X.transpose(1, 2)
+    b = torch.rand((X.size(0), X.size(1))).cuda()
+    idx = b.sort(0)[1]
+    adx = torch.range(0, X.size(1) - 1).long()
+    X = X[idx, adx[None, :]].transpose(1, 2)
+    return X

cortex_DIM/nn_modules/convnet.py

@@ -0,0 +1,352 @@
+'''Convnet encoder module.
+
+'''
+
+import torch
+import torch.nn as nn
+
+#from cortex.built_ins.networks.utils import get_nonlinearity
+
+from cortex_DIM.nn_modules.misc import Fold, Unfold, View
+
+
+def infer_conv_size(w, k, s, p):
+    '''Infers the next size after convolution.
+
+    Args:
+        w: Input size.
+        k: Kernel size.
+        s: Stride.
+        p: Padding.
+
+    Returns:
+        int: Output size.
+
+    '''
+    x = (w - k + 2 * p) // s + 1
+    return x
+
+
+class Convnet(nn.Module):
+    '''Basic convnet convenience class.
+
+    Attributes:
+        conv_layers: nn.Sequential of nn.Conv2d layers with batch norm,
+            dropout, nonlinearity.
+        fc_layers: nn.Sequential of nn.Linear layers with batch norm,
+            dropout, nonlinearity.
+        reshape: Simple reshape layer.
+        conv_shape: Shape of the convolutional output.
+
+    '''
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.create_layers(*args, **kwargs)
+
+    def create_layers(self, shape, conv_args=None, fc_args=None):
+        '''Creates layers
+
+        conv_args are in format (dim_h, f_size, stride, pad, batch_norm, dropout, nonlinearity, pool)
+        fc_args are in format (dim_h, batch_norm, dropout, nonlinearity)
+
+        Args:
+            shape: Shape of input.
+            conv_args: List of tuple of convolutional arguments.
+            fc_args: List of tuple of fully-connected arguments.
+        '''
+
+        self.conv_layers, self.conv_shape = self.create_conv_layers(shape, conv_args)
+
+        dim_x, dim_y, dim_out = self.conv_shape
+        dim_r = dim_x * dim_y * dim_out
+        self.reshape = View(-1, dim_r)
+        self.fc_layers, _ = self.create_linear_layers(dim_r, fc_args)
+
+    def create_conv_layers(self, shape, conv_args):
+        '''Creates a set of convolutional layers.
+
+        Args:
+            shape: Input shape.
+            conv_args: List of tuple of convolutional arguments.
+
+        Returns:
+            nn.Sequential: a sequence of convolutional layers.
+
+        '''
+
+        conv_layers = nn.Sequential()
+        conv_args = conv_args or []
+
+        dim_x, dim_y, dim_in = shape
+
+        for i, (dim_out, f, s, p, batch_norm, dropout, nonlinearity, pool) in enumerate(conv_args):
+            name = '({}/{})_{}'.format(dim_in, dim_out, i + 1)
+            conv_block = nn.Sequential()
+
+            if dim_out is not None:
+                conv = nn.Conv2d(dim_in, dim_out, kernel_size=f, stride=s, padding=p, bias=not(batch_norm))
+                conv_block.add_module(name + 'conv', conv)
+                dim_x, dim_y = self.next_size(dim_x, dim_y, f, s, p)
+            else:
+                dim_out = dim_in
+
+            if dropout:
+                conv_block.add_module(name + 'do', nn.Dropout2d(p=dropout))
+            if batch_norm:
+                bn = nn.BatchNorm2d(dim_out)
+                conv_block.add_module(name + 'bn', bn)
+
+            if nonlinearity:
+                nonlinearity = get_nonlinearity(nonlinearity)
+                conv_block.add_module(nonlinearity.__class__.__name__, nonlinearity)
+
+            if pool:
+                (pool_type, kernel, stride) = pool
+                Pool = getattr(nn, pool_type)
+                conv_block.add_module(name + 'pool', Pool(kernel_size=kernel, stride=stride))
+                dim_x, dim_y = self.next_size(dim_x, dim_y, kernel, stride, 0)
+
+            conv_layers.add_module(name, conv_block)
+
+            dim_in = dim_out
+
+        dim_out = dim_in
+
+        return conv_layers, (dim_x, dim_y, dim_out)
+
+    def create_linear_layers(self, dim_in, fc_args):
+        '''
+
+        Args:
+            dim_in: Number of input units.
+            fc_args: List of tuple of fully-connected arguments.
+
+        Returns:
+            nn.Sequential.
+
+        '''
+
+        fc_layers = nn.Sequential()
+        fc_args = fc_args or []
+
+        for i, (dim_out, batch_norm, dropout, nonlinearity) in enumerate(fc_args):
+            name = '({}/{})_{}'.format(dim_in, dim_out, i + 1)
+            fc_block = nn.Sequential()
+
+            if dim_out is not None:
+                fc_block.add_module(name + 'fc', nn.Linear(dim_in, dim_out))
+            else:
+                dim_out = dim_in
+
+            if dropout:
+                fc_block.add_module(name + 'do', nn.Dropout(p=dropout))
+            if batch_norm:
+                bn = nn.BatchNorm1d(dim_out)
+                fc_block.add_module(name + 'bn', bn)
+
+            if nonlinearity:
+                nonlinearity = get_nonlinearity(nonlinearity)
+                fc_block.add_module(nonlinearity.__class__.__name__, nonlinearity)
+
+            fc_layers.add_module(name, fc_block)
+
+            dim_in = dim_out
+
+        return fc_layers, dim_in
+
+    def next_size(self, dim_x, dim_y, k, s, p):
+        '''Infers the next size of a convolutional layer.
+
+        Args:
+            dim_x: First dimension.
+            dim_y: Second dimension.
+            k: Kernel size.
+            s: Stride.
+            p: Padding.
+
+        Returns:
+            (int, int): (First output dimension, Second output dimension)
+
+        '''
+        if isinstance(k, int):
+            kx, ky = (k, k)
+        else:
+            kx, ky = k
+
+        if isinstance(s, int):
+            sx, sy = (s, s)
+        else:
+            sx, sy = s
+
+        if isinstance(p, int):
+            px, py = (p, p)
+        else:
+            px, py = p
+        return (infer_conv_size(dim_x, kx, sx, px),
+                infer_conv_size(dim_y, ky, sy, py))
+
+    def forward(self, x: torch.Tensor, return_full_list=False):
+        '''Forward pass
+
+        Args:
+            x: Input.
+            return_full_list: Optional, returns all layer outputs.
+
+        Returns:
+            torch.Tensor or list of torch.Tensor.
+
+        '''
+        if return_full_list:
+            conv_out = []
+            for conv_layer in self.conv_layers:
+                x = conv_layer(x)
+                conv_out.append(x)
+        else:
+            conv_out = self.conv_layers(x)
+            x = conv_out
+
+        x = self.reshape(x)
+
+        if return_full_list:
+            fc_out = []
+            for fc_layer in self.fc_layers:
+                x = fc_layer(x)
+                fc_out.append(x)
+        else:
+            fc_out = self.fc_layers(x)
+
+        return conv_out, fc_out
+
+
+class FoldedConvnet(Convnet):
+    '''Convnet with strided crop input.
+
+    '''
+
+    def create_layers(self, shape, crop_size=8, conv_args=None, fc_args=None):
+        '''Creates layers
+
+        conv_args are in format (dim_h, f_size, stride, pad, batch_norm, dropout, nonlinearity, pool)
+        fc_args are in format (dim_h, batch_norm, dropout, nonlinearity)
+
+        Args:
+            shape: Shape of input.
+            crop_size: Size of crops
+            conv_args: List of tuple of convolutional arguments.
+            fc_args: List of tuple of fully-connected arguments.
+        '''
+
+        self.crop_size = crop_size
+
+        dim_x, dim_y, dim_in = shape
+        if dim_x != dim_y:
+            raise ValueError('x and y dimensions must be the same to use Folded encoders.')
+
+        self.final_size = 2 * (dim_x // self.crop_size) - 1
+
+        self.unfold = Unfold(dim_x, self.crop_size)
+        self.refold = Fold(dim_x, self.crop_size)
+
+        shape = (self.crop_size, self.crop_size, dim_in)
+
+        self.conv_layers, self.conv_shape = self.create_conv_layers(shape, conv_args)
+
+        dim_x, dim_y, dim_out = self.conv_shape
+        dim_r = dim_x * dim_y * dim_out
+        self.reshape = View(-1, dim_r)
+        self.fc_layers, _ = self.create_linear_layers(dim_r, fc_args)
+
+    def create_conv_layers(self, shape, conv_args):
+        '''Creates a set of convolutional layers.
+
+        Args:
+            shape: Input shape.
+            conv_args: List of tuple of convolutional arguments.
+
+        Returns:
+            nn.Sequential: A sequence of convolutional layers.
+
+        '''
+
+        conv_layers = nn.Sequential()
+        conv_args = conv_args or []
+        dim_x, dim_y, dim_in = shape
+
+        for i, (dim_out, f, s, p, batch_norm, dropout, nonlinearity, pool) in enumerate(conv_args):
+            name = '({}/{})_{}'.format(dim_in, dim_out, i + 1)
+            conv_block = nn.Sequential()
+
+            if dim_out is not None:
+                conv = nn.Conv2d(dim_in, dim_out, kernel_size=f, stride=s, padding=p, bias=not(batch_norm))
+                conv_block.add_module(name + 'conv', conv)
+                dim_x, dim_y = self.next_size(dim_x, dim_y, f, s, p)
+            else:
+                dim_out = dim_in
+
+            if dropout:
+                conv_block.add_module(name + 'do', nn.Dropout2d(p=dropout))
+            if batch_norm:
+                bn = nn.BatchNorm2d(dim_out)
+                conv_block.add_module(name + 'bn', bn)
+
+            if nonlinearity:
+                nonlinearity = get_nonlinearity(nonlinearity)
+                conv_block.add_module(nonlinearity.__class__.__name__, nonlinearity)
+
+            if pool:
+                (pool_type, kernel, stride) = pool
+                Pool = getattr(nn, pool_type)
+                conv_block.add_module('pool', Pool(kernel_size=kernel, stride=stride))
+                dim_x, dim_y = self.next_size(dim_x, dim_y, kernel, stride, 0)
+
+            conv_layers.add_module(name, conv_block)
+
+            dim_in = dim_out
+
+            if dim_x != dim_y:
+                raise ValueError('dim_x and dim_y do not match.')
+
+            if dim_x == 1:
+                dim_x = self.final_size
+                dim_y = self.final_size
+
+        dim_out = dim_in
+
+        return conv_layers, (dim_x, dim_y, dim_out)
+
+    def forward(self, x: torch.Tensor, return_full_list=False):
+        '''Forward pass
+
+        Args:
+            x: Input.
+            return_full_list: Optional, returns all layer outputs.
+
+        Returns:
+            torch.Tensor or list of torch.Tensor.
+
+        '''
+
+        x = self.unfold(x)
+
+        conv_out = []
+        for conv_layer in self.conv_layers:
+            x = conv_layer(x)
+            if x.size(2) == 1:
+                x = self.refold(x)
+            conv_out.append(x)
+
+        x = self.reshape(x)
+
+        if return_full_list:
+            fc_out = []
+            for fc_layer in self.fc_layers:
+                x = fc_layer(x)
+                fc_out.append(x)
+        else:
+            fc_out = self.fc_layers(x)
+
+        if not return_full_list:
+            conv_out = conv_out[-1]
+
+        return conv_out, fc_out

cortex_DIM/nn_modules/encoder.py

@@ -0,0 +1,96 @@
+'''Basic cortex_DIM encoder.
+
+'''
+
+import torch
+
+from cortex_DIM.nn_modules.convnet import Convnet, FoldedConvnet
+from cortex_DIM.nn_modules.resnet import ResNet, FoldedResNet
+
+
+def create_encoder(Module):
+    class Encoder(Module):
+        '''Encoder used for cortex_DIM.
+
+        '''
+
+        def __init__(self, *args, local_idx=None, multi_idx=None, conv_idx=None, fc_idx=None, **kwargs):
+            '''
+
+            Args:
+                args: Arguments for parent class.
+                local_idx: Index in list of convolutional layers for local features.
+                multi_idx: Index in list of convolutional layers for multiple globals.
+                conv_idx: Index in list of convolutional layers for intermediate features.
+                fc_idx: Index in list of fully-connected layers for intermediate features.
+                kwargs: Keyword arguments for the parent class.
+            '''
+
+            super().__init__(*args, **kwargs)
+
+            if local_idx is None:
+                raise ValueError('`local_idx` must be set')
+
+            conv_idx = conv_idx or local_idx
+
+            self.local_idx = local_idx
+            self.multi_idx = multi_idx
+            self.conv_idx = conv_idx
+            self.fc_idx = fc_idx
+
+        def forward(self, x: torch.Tensor):
+            '''
+
+            Args:
+                x: Input tensor.
+
+            Returns:
+                local_out, multi_out, hidden_out, global_out
+
+            '''
+
+            outs = super().forward(x, return_full_list=True)
+            if len(outs) == 2:
+                conv_out, fc_out = outs
+            else:
+                conv_before_out, res_out, conv_after_out, fc_out = outs
+                conv_out = conv_before_out + res_out + conv_after_out
+
+            local_out = conv_out[self.local_idx]
+
+            if self.multi_idx is not None:
+                multi_out = conv_out[self.multi_idx]
+            else:
+                multi_out = None
+
+            if len(fc_out) > 0:
+                if self.fc_idx is not None:
+                    hidden_out = fc_out[self.fc_idx]
+                else:
+                    hidden_out = None
+                global_out = fc_out[-1]
+            else:
+                hidden_out = None
+                global_out = None
+
+            conv_out = conv_out[self.conv_idx]
+
+            return local_out, conv_out, multi_out, hidden_out, global_out
+
+    return Encoder
+
+
+class ConvnetEncoder(create_encoder(Convnet)):
+    pass
+
+
+class FoldedConvnetEncoder(create_encoder(FoldedConvnet)):
+    pass
+
+
+class ResnetEncoder(create_encoder(ResNet)):
+    pass
+
+
+class FoldedResnetEncoder(create_encoder(FoldedResNet)):
+    pass

cortex_DIM/nn_modules/mi_networks.py

@@ -0,0 +1,106 @@
+"""Module for networks used for computing MI.
+
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from cortex_DIM.nn_modules.misc import Permute
+
+
+class MIFCNet(nn.Module):
+    """Simple custom network for computing MI.
+
+    """
+    def __init__(self, n_input, n_units):
+        """
+
+        Args:
+            n_input: Number of input units.
+            n_units: Number of output units.
+        """
+        super().__init__()
+
+        assert(n_units >= n_input)
+
+        self.linear_shortcut = nn.Linear(n_input, n_units)
+        self.block_nonlinear = nn.Sequential(
+            nn.Linear(n_input, n_units),
+            nn.BatchNorm1d(n_units),
+            nn.ReLU(),
+            nn.Linear(n_units, n_units)
+        )
+
+        # initialize the initial projection to a sort of noisy copy
+        eye_mask = np.zeros((n_units, n_input), dtype=np.uint8)
+        for i in range(n_input):
+            eye_mask[i, i] = 1
+
+        self.linear_shortcut.weight.data.uniform_(-0.01, 0.01)
+        self.linear_shortcut.weight.data.masked_fill_(torch.tensor(eye_mask), 1.)
+
+    def forward(self, x):
+        """
+
+        Args:
+            x: Input tensor.
+
+        Returns:
+            torch.Tensor: network output.
+
+        """
+        h = self.block_nonlinear(x) + self.linear_shortcut(x)
+        return h
+
+
+class MI1x1ConvNet(nn.Module):
+    """Simple custorm 1x1 convnet.
+
+    """
+    def __init__(self, n_input, n_units):
+        """
+
+        Args:
+            n_input: Number of input units.
+            n_units: Number of output units.
+        """
+
+        super().__init__()
+
+        self.block_nonlinear = nn.Sequential(
+            nn.Conv1d(n_input, n_units, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.BatchNorm1d(n_units),
+            nn.ReLU(),
+            nn.Conv1d(n_units, n_units, kernel_size=1, stride=1, padding=0, bias=True),
+        )
+
+        self.block_ln = nn.Sequential(
+            Permute(0, 2, 1),
+            nn.LayerNorm(n_units),
+            Permute(0, 2, 1)
+        )
+
+        self.linear_shortcut = nn.Conv1d(n_input, n_units, kernel_size=1,
+                                         stride=1, padding=0, bias=False)
+
+        # initialize shortcut to be like identity (if possible)
+        if n_units >= n_input:
+            eye_mask = np.zeros((n_units, n_input, 1), dtype=np.uint8)
+            for i in range(n_input):
+                eye_mask[i, i, 0] = 1
+            self.linear_shortcut.weight.data.uniform_(-0.01, 0.01)
+            self.linear_shortcut.weight.data.masked_fill_(torch.tensor(eye_mask), 1.)
+
+    def forward(self, x):
+        """
+
+            Args:
+                x: Input tensor.
+
+            Returns:
+                torch.Tensor: network output.
+
+        """
+        h = self.block_ln(self.block_nonlinear(x) + self.linear_shortcut(x))
+        return h

cortex_DIM/nn_modules/misc.py

@@ -0,0 +1,130 @@
+'''Various miscellaneous modules
+
+'''
+
+import torch
+
+
+class View(torch.nn.Module):
+    """Basic reshape module.
+
+    """
+    def __init__(self, *shape):
+        """
+
+        Args:
+            *shape: Input shape.
+        """
+        super().__init__()
+        self.shape = shape
+
+    def forward(self, input):
+        """Reshapes tensor.
+
+        Args:
+            input: Input tensor.
+
+        Returns:
+            torch.Tensor: Flattened tensor.
+
+        """
+        return input.view(*self.shape)
+
+
+class Unfold(torch.nn.Module):
+    """Module for unfolding tensor.
+
+    Performs strided crops on 2d (image) tensors. Stride is assumed to be half the crop size.
+
+    """
+    def __init__(self, img_size, fold_size):
+        """
+
+        Args:
+            img_size: Input size.
+            fold_size: Crop size.
+        """
+        super().__init__()
+
+        fold_stride = fold_size // 2
+        self.fold_size = fold_size
+        self.fold_stride = fold_stride
+        self.n_locs = 2 * (img_size // fold_size) - 1
+        self.unfold = torch.nn.Unfold((self.fold_size, self.fold_size),
+                                      stride=(self.fold_stride, self.fold_stride))
+
+    def forward(self, x):
+        """Unfolds tensor.
+
+        Args:
+            x: Input tensor.
+
+        Returns:
+            torch.Tensor: Unfolded tensor.
+
+        """
+        N = x.size(0)
+        x = self.unfold(x).reshape(N, -1, self.fold_size, self.fold_size, self.n_locs * self.n_locs)\
+            .permute(0, 4, 1, 2, 3)\
+            .reshape(N * self.n_locs * self.n_locs, -1, self.fold_size, self.fold_size)
+        return x
+
+
+class Fold(torch.nn.Module):
+    """Module (re)folding tensor.
+
+    Undoes the strided crops above. Works only on 1x1.
+
+    """
+    def __init__(self, img_size, fold_size):
+        """
+
+        Args:
+            img_size: Images size.
+            fold_size: Crop size.
+        """
+        super().__init__()
+        self.n_locs = 2 * (img_size // fold_size) - 1
+
+    def forward(self, x):
+        """(Re)folds tensor.
+
+        Args:
+            x: Input tensor.
+
+        Returns:
+            torch.Tensor: Refolded tensor.
+
+        """
+        dim_c, dim_x, dim_y = x.size()[1:]
+        x = x.reshape(-1, self.n_locs * self.n_locs, dim_c, dim_x * dim_y)
+        x = x.reshape(-1, self.n_locs * self.n_locs, dim_c, dim_x * dim_y)\
+            .permute(0, 2, 3, 1)\
+            .reshape(-1, dim_c * dim_x * dim_y, self.n_locs, self.n_locs).contiguous()
+        return x
+
+
+class Permute(torch.nn.Module):
+    """Module for permuting axes.
+
+    """
+    def __init__(self, *perm):
+        """
+
+        Args:
+            *perm: Permute axes.
+        """
+        super().__init__()
+        self.perm = perm
+
+    def forward(self, input):
+        """Permutes axes of tensor.
+
+        Args:
+            input: Input tensor.
+
+        Returns:
+            torch.Tensor: permuted tensor.
+
+        """
+        return input.permute(*self.perm)

cortex_DIM/nn_modules/resnet.py

@@ -0,0 +1,297 @@
+'''Module for making resnet encoders.
+
+'''
+
+import torch
+import torch.nn as nn
+
+from cortex_DIM.nn_modules.convnet import Convnet
+from cortex_DIM.nn_modules.misc import Fold, Unfold, View
+
+
+_nonlin_idx = 6
+
+
+class ResBlock(Convnet):
+    '''Residual block for ResNet
+
+    '''
+
+    def create_layers(self, shape, conv_args=None):
+        '''Creates layers
+
+        Args:
+            shape: Shape of input.
+            conv_args: Layer arguments for block.
+        '''
+
+        # Move nonlinearity to a separate step for residual.
+        final_nonlin = conv_args[-1][_nonlin_idx]
+        conv_args[-1] = list(conv_args[-1])
+        conv_args[-1][_nonlin_idx] = None
+        conv_args.append((None, 0, 0, 0, False, False, final_nonlin, None))
+
+        super().create_layers(shape, conv_args=conv_args)
+
+        if self.conv_shape != shape:
+            dim_x, dim_y, dim_in = shape
+            dim_x_, dim_y_, dim_out = self.conv_shape
+            stride = dim_x // dim_x_
+            next_x, _ = self.next_size(dim_x, dim_y, 1, stride, 0)
+            assert next_x == dim_x_, (self.conv_shape, shape)
+
+            self.downsample = nn.Sequential(
+                nn.Conv2d(dim_in, dim_out, kernel_size=1, stride=stride, padding=0, bias=False),
+                nn.BatchNorm2d(dim_out),
+            )
+        else:
+            self.downsample = None
+
+    def forward(self, x: torch.Tensor):
+        '''Forward pass
+
+        Args:
+            x: Input.
+
+        Returns:
+            torch.Tensor or list of torch.Tensor.
+
+        '''
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        else:
+            residual = x
+
+        x = self.conv_layers[-1](self.conv_layers[:-1](x) + residual)
+
+        return x
+
+
+class ResNet(Convnet):
+    def create_layers(self, shape, conv_before_args=None, res_args=None, conv_after_args=None, fc_args=None):
+        '''Creates layers
+
+        Args:
+            shape: Shape of the input.
+            conv_before_args: Arguments for convolutional layers before residuals.
+            res_args: Residual args.
+            conv_after_args: Arguments for convolutional layers after residuals.
+            fc_args: Fully-connected arguments.
+
+        '''
+
+        dim_x, dim_y, dim_in = shape
+        shape = (dim_x, dim_y, dim_in)
+        self.conv_before_layers, self.conv_before_shape = self.create_conv_layers(shape, conv_before_args)
+        self.res_layers, self.res_shape = self.create_res_layers(self.conv_before_shape, res_args)
+        self.conv_after_layers, self.conv_after_shape = self.create_conv_layers(self.res_shape, conv_after_args)
+
+        dim_x, dim_y, dim_out = self.conv_after_shape
+        dim_r = dim_x * dim_y * dim_out
+        self.reshape = View(-1, dim_r)
+        self.fc_layers, _ = self.create_linear_layers(dim_r, fc_args)
+
+    def create_res_layers(self, shape, block_args=None):
+        '''Creates a set of residual blocks.
+
+        Args:
+            shape: input shape.
+            block_args: Arguments for blocks.
+
+        Returns:
+            nn.Sequential: sequence of residual blocks.
+
+        '''
+
+        res_layers = nn.Sequential()
+        block_args = block_args or []
+
+        for i, (conv_args, n_blocks) in enumerate(block_args):
+            block = ResBlock(shape, conv_args=conv_args)
+            res_layers.add_module('block_{}_0'.format(i), block)
+
+            for j in range(1, n_blocks):
+                shape = block.conv_shape
+                block = ResBlock(shape, conv_args=conv_args)
+                res_layers.add_module('block_{}_{}'.format(i, j), block)
+            shape = block.conv_shape
+
+        return res_layers, shape
+
+    def forward(self, x: torch.Tensor, return_full_list=False):
+        '''Forward pass
+
+        Args:
+            x: Input.
+            return_full_list: Optional, returns all layer outputs.
+
+        Returns:
+            torch.Tensor or list of torch.Tensor.
+
+        '''
+
+        if return_full_list:
+            conv_before_out = []
+            for conv_layer in self.conv_before_layers:
+                x = conv_layer(x)
+                conv_before_out.append(x)
+        else:
+            conv_before_out = self.conv_layers(x)
+            x = conv_before_out
+
+        if return_full_list:
+            res_out = []
+            for res_layer in self.res_layers:
+                x = res_layer(x)
+                res_out.append(x)
+        else:
+            res_out = self.res_layers(x)
+            x = res_out
+
+        if return_full_list:
+            conv_after_out = []
+            for conv_layer in self.conv_after_layers:
+                x = conv_layer(x)
+                conv_after_out.append(x)
+        else:
+            conv_after_out = self.conv_after_layers(x)
+            x = conv_after_out
+
+        x = self.reshape(x)
+
+        if return_full_list:
+            fc_out = []
+            for fc_layer in self.fc_layers:
+                x = fc_layer(x)
+                fc_out.append(x)
+        else:
+            fc_out = self.fc_layers(x)
+
+        return conv_before_out, res_out, conv_after_out, fc_out
+
+
+class FoldedResNet(ResNet):
+    '''Resnet with strided crop input.
+
+    '''
+
+    def create_layers(self, shape, crop_size=8, conv_before_args=None, res_args=None,
+                      conv_after_args=None, fc_args=None):
+        '''Creates layers
+
+        Args:
+            shape: Shape of the input.
+            crop_size: Size of the crops.
+            conv_before_args: Arguments for convolutional layers before residuals.
+            res_args: Residual args.
+            conv_after_args: Arguments for convolutional layers after residuals.
+            fc_args: Fully-connected arguments.
+
+        '''
+        self.crop_size = crop_size
+
+        dim_x, dim_y, dim_in = shape
+        self.final_size = 2 * (dim_x // self.crop_size) - 1
+
+        self.unfold = Unfold(dim_x, self.crop_size)
+        self.refold = Fold(dim_x, self.crop_size)
+
+        shape = (self.crop_size, self.crop_size, dim_in)
+        self.conv_before_layers, self.conv_before_shape = self.create_conv_layers(shape, conv_before_args)
+
+        self.res_layers, self.res_shape = self.create_res_layers(self.conv_before_shape, res_args)
+        self.conv_after_layers, self.conv_after_shape = self.create_conv_layers(self.res_shape, conv_after_args)
+        self.conv_after_shape = self.res_shape
+
+        dim_x, dim_y, dim_out = self.conv_after_shape
+        dim_r = dim_x * dim_y * dim_out
+        self.reshape = View(-1, dim_r)
+        self.fc_layers, _ = self.create_linear_layers(dim_r, fc_args)
+
+    def create_res_layers(self, shape, block_args=None):
+        '''Creates a set of residual blocks.
+
+        Args:
+            shape: input shape.
+            block_args: Arguments for blocks.
+
+        Returns:
+            nn.Sequential: sequence of residual blocks.
+
+        '''
+
+        res_layers = nn.Sequential()
+        block_args = block_args or []
+
+        for i, (conv_args, n_blocks) in enumerate(block_args):
+            block = ResBlock(shape, conv_args=conv_args)
+            res_layers.add_module('block_{}_0'.format(i), block)
+
+            for j in range(1, n_blocks):
+                shape = block.conv_shape
+                block = ResBlock(shape, conv_args=conv_args)
+                res_layers.add_module('block_{}_{}'.format(i, j), block)
+            shape = block.conv_shape
+            dim_x, dim_y = shape[:2]
+
+            if dim_x != dim_y:
+                raise ValueError('dim_x and dim_y do not match.')
+
+            if dim_x == 1:
+                shape = (self.final_size, self.final_size, shape[2])
+
+        return res_layers, shape
+
+    def forward(self, x: torch.Tensor, return_full_list=False):
+        '''Forward pass
+
+        Args:
+            x: Input.
+            return_full_list: Optional, returns all layer outputs.
+
+        Returns:
+            torch.Tensor or list of torch.Tensor.
+
+        '''
+        x = self.unfold(x)
+
+        conv_before_out = []
+        for conv_layer in self.conv_before_layers:
+            x = conv_layer(x)
+            if x.size(2) == 1:
+                x = self.refold(x)
+            conv_before_out.append(x)
+
+        res_out = []
+        for res_layer in self.res_layers:
+            x = res_layer(x)
+            res_out.append(x)
+
+        if x.size(2) == 1:
+            x = self.refold(x)
+            res_out[-1] = x
+
+        conv_after_out = []
+        for conv_layer in self.conv_after_layers:
+            x = conv_layer(x)
+            if x.size(2) == 1:
+                x = self.refold(x)
+            conv_after_out.append(x)
+
+        x = self.reshape(x)
+
+        if return_full_list:
+            fc_out = []
+            for fc_layer in self.fc_layers:
+                x = fc_layer(x)
+                fc_out.append(x)
+        else:
+            fc_out = self.fc_layers(x)
+
+        if not return_full_list:
+            conv_before_out = conv_before_out[-1]
+            res_out = res_out[-1]
+            conv_after_out = conv_after_out[-1]
+
+        return conv_before_out, res_out, conv_after_out, fc_out

gin.py

@@ -0,0 +1,188 @@
+from sklearn import preprocessing
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import accuracy_score
+from sklearn.model_selection import GridSearchCV, KFold, StratifiedKFold
+from sklearn.model_selection import cross_val_score
+from sklearn.svm import SVC, LinearSVC
+from torch.nn import Sequential, Linear, ReLU
+from torch_geometric.data import DataLoader
+from torch_geometric.datasets import TUDataset
+from torch_geometric.nn import GINConv, global_add_pool
+from tqdm import tqdm
+import numpy as np
+import os.path as osp
+import sys
+import torch
+import torch.nn.functional as F
+
+
+class Encoder(torch.nn.Module):
+    def __init__(self, num_features, dim, num_gc_layers):
+        super(Encoder, self).__init__()
+
+        # num_features = dataset.num_features
+        # dim = 32
+        self.num_gc_layers = num_gc_layers
+
+        # self.nns = []
+        self.convs = torch.nn.ModuleList()
+        self.bns = torch.nn.ModuleList()
+
+        for i in range(num_gc_layers):
+
+            if i:
+                nn = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))
+            else:
+                nn = Sequential(Linear(num_features, dim), ReLU(), Linear(dim, dim))
+            conv = GINConv(nn)
+            bn = torch.nn.BatchNorm1d(dim)
+
+            self.convs.append(conv)
+            self.bns.append(bn)
+
+
+    def forward(self, x, edge_index, batch):
+        if x is None:
+            x = torch.ones((batch.shape[0], 1)).to(device)
+
+        xs = []
+        for i in range(self.num_gc_layers):
+
+            x = F.relu(self.convs[i](x, edge_index))
+            x = self.bns[i](x)
+            xs.append(x)
+            # if i == 2:
+                # feature_map = x2
+
+        xpool = [global_add_pool(x, batch) for x in xs]
+        x = torch.cat(xpool, 1)
+        return x, torch.cat(xs, 1)
+
+    def get_embeddings(self, loader):
+
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        ret = []
+        y = []
+        with torch.no_grad():
+            for data in loader:
+                data.to(device)
+                x, edge_index, batch = data.x, data.edge_index, data.batch
+                if x is None:
+                    x = torch.ones((batch.shape[0],1)).to(device)
+                x, _ = self.forward(x, edge_index, batch)
+                ret.append(x.cpu().numpy())
+                # y.append(data.aid)
+        ret = np.concatenate(ret, 0)
+        # y = np.concatenate(y, 0)
+        return ret
+        # return ret, y
+
+class Net(torch.nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+
+        try:
+            num_features = dataset.num_features
+        except:
+            num_features = 1
+        dim = 32
+
+        self.encoder = Encoder(num_features, dim)
+
+        self.fc1 = Linear(dim*5, dim)
+        self.fc2 = Linear(dim, dataset.num_classes)
+
+    def forward(self, x, edge_index, batch):
+        if x is None:
+            x = torch.ones(batch.shape[0]).to(device)
+
+        x, _ = self.encoder(x, edge_index, batch)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, p=0.5, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=-1)
+
+def train(epoch):
+    model.train()
+
+    if epoch == 51:
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = 0.5 * param_group['lr']
+
+    loss_all = 0
+    for data in train_loader:
+        data = data.to(device)
+        optimizer.zero_grad()
+        # print(data.x.shape)
+        # [ num_nodes x num_node_labels ]
+        # print(data.edge_index.shape)
+        #  [2 x num_edges ]
+        # print(data.batch.shape)
+        # [ num_nodes ]
+        output = model(data.x, data.edge_index, data.batch)
+        loss = F.nll_loss(output, data.y)
+        loss.backward()
+        loss_all += loss.item() * data.num_graphs
+        optimizer.step()
+
+    return loss_all / len(train_dataset)
+
+def test(loader):
+    model.eval()
+
+    correct = 0
+    for data in loader:
+        data = data.to(device)
+        output = model(data.x, data.edge_index, data.batch)
+        pred = output.max(dim=1)[1]
+        correct += pred.eq(data.y).sum().item()
+    return correct / len(loader.dataset)
+
+
+if __name__ == '__main__':
+    for percentage in [ 1.]:
+        for DS in [sys.argv[1]]:
+            if 'REDDIT' in DS:
+                epochs = 200
+            else:
+                epochs = 100
+            path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', DS)
+            accuracies = [[] for i in range(epochs)]
+            #kf = StratifiedKFold(n_splits=10, shuffle=True, random_state=None)
+            dataset = TUDataset(path, name=DS) #.shuffle()
+            num_graphs = len(dataset)
+            print('Number of graphs', len(dataset))
+            dataset = dataset[:int(num_graphs * percentage)]
+            dataset = dataset.shuffle()
+
+            kf = KFold(n_splits=10, shuffle=True, random_state=None)
+            for train_index, test_index in kf.split(dataset):
+
+                # x_train, x_test = x[train_index], x[test_index]
+                # y_train, y_test = y[train_index], y[test_index]
+                train_dataset = [dataset[int(i)] for i in list(train_index)]
+                test_dataset = [dataset[int(i)] for i in list(test_index)]
+                print('len(train_dataset)', len(train_dataset))
+                print('len(test_dataset)', len(test_dataset))
+
+                train_loader = DataLoader(train_dataset, batch_size=128)
+                test_loader = DataLoader(test_dataset, batch_size=128)
+                # print('train', len(train_loader))
+                # print('test', len(test_loader))
+
+                device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+                model = Net().to(device)
+                optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+                for epoch in range(1, epochs+1):
+                    train_loss = train(epoch)
+                    train_acc = test(train_loader)
+                    test_acc = test(test_loader)
+                    accuracies[epoch-1].append(test_acc)
+                    tqdm.write('Epoch: {:03d}, Train Loss: {:.7f}, '
+                          'Train Acc: {:.7f}, Test Acc: {:.7f}'.format(epoch, train_loss,
+                                                                       train_acc, test_acc))
+            tmp = np.mean(accuracies, axis=1)
+            print(percentage, DS, np.argmax(tmp), np.max(tmp), np.std(accuracies[np.argmax(tmp)]))
+            input()

model.py

@@ -0,0 +1,66 @@
+from cortex_DIM.nn_modules.mi_networks import MIFCNet, MI1x1ConvNet
+from torch import optim
+from torch.autograd import Variable
+import json
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class GlobalDiscriminator(nn.Module):
+    def __init__(self, args, input_dim):
+        super().__init__()
+        
+        self.l0 = nn.Linear(32, 32)
+        self.l1 = nn.Linear(32, 32)
+
+        self.l2 = nn.Linear(512, 1)
+    def forward(self, y, M, data):
+
+        adj = Variable(data['adj'].float(), requires_grad=False).cuda()
+        # h0 = Variable(data['feats'].float()).cuda()
+        batch_num_nodes = data['num_nodes'].int().numpy()
+        M, _ = self.encoder(M, adj, batch_num_nodes)
+        # h = F.relu(self.c0(M))
+        # h = self.c1(h)
+        # h = h.view(y.shape[0], -1)
+        h = torch.cat((y, M), dim=1)
+        h = F.relu(self.l0(h))
+        h = F.relu(self.l1(h))
+        return self.l2(h)
+
+class PriorDiscriminator(nn.Module):
+    def __init__(self, input_dim):
+        super().__init__()
+        self.l0 = nn.Linear(input_dim, input_dim)
+        self.l1 = nn.Linear(input_dim, input_dim)
+        self.l2 = nn.Linear(input_dim, 1)
+
+    def forward(self, x):
+        h = F.relu(self.l0(x))
+        h = F.relu(self.l1(h))
+        return torch.sigmoid(self.l2(h))
+
+class FF(nn.Module):
+    def __init__(self, input_dim):
+        super().__init__()
+        # self.c0 = nn.Conv1d(input_dim, 512, kernel_size=1)
+        # self.c1 = nn.Conv1d(512, 512, kernel_size=1)
+        # self.c2 = nn.Conv1d(512, 1, kernel_size=1)
+        self.block = nn.Sequential(
+            nn.Linear(input_dim, input_dim),
+            nn.ReLU(),
+            nn.Linear(input_dim, input_dim),
+            nn.ReLU(),
+            nn.Linear(input_dim, input_dim),
+            nn.ReLU()
+        )
+        self.linear_shortcut = nn.Linear(input_dim, input_dim)
+        # self.c0 = nn.Conv1d(input_dim, 512, kernel_size=1, stride=1, padding=0)
+        # self.c1 = nn.Conv1d(512, 512, kernel_size=1, stride=1, padding=0)
+        # self.c2 = nn.Conv1d(512, 1, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        return self.block(x) + self.linear_shortcut(x)
+

requirements.txt

@@ -0,0 +1,8 @@
+numpy==1.25.0
+scikit-learn==1.2.2
+shapely==2.0.1
+statistics
+collection
+torch==2.0.1
+opencv-python==4.8.0.74
+torch-geometric==2.4.0

utils/clean_data.py

@@ -0,0 +1,192 @@
+import numpy as np
+from collections import defaultdict
+import shapely
+import shapely.wkt
+from shapely.geometry import LineString, MultiLineString, Polygon, Point, MultiPoint
+from shapely.prepared import prep
+
+
+def list_duplicates(seq):
+    tally = defaultdict(list)
+    for i,item in enumerate(seq):
+        tally[item].append(i)
+    return ((key,locs) for key,locs in tally.items() if len(locs)>1)
+
+
+# clean wall and wkt
+# wall = ((IOOI))
+# wkt = 'POLYGON ((x0 y0, x1 y1, x2 y2, x3 y3, x0 y0))'
+def read_wall_wkt(wall, wkt):
+    # wall to list
+    wall_l = wall.split("), (")[0]
+    wall_l = wall_l.split("((")[1]
+    wall_l = wall_l.split("))")[0]
+    wall_c = [*wall_l]
+
+    #clean wkt
+    wkt_l = wkt.split("((")[1]
+    wkt_l = wkt_l.split("))")[0]
+    wkt_l = wkt_l.split("), (")
+
+    if len(wkt_l) == 1:
+        wkt_c = wkt
+    else:
+        wkt_c = "POLYGON ((" + wkt_l[0] + "))"
+
+    wkt_c = wkt_c.split("((")[1]
+    wkt_c = wkt_c.split("))")[0]  
+    wkt_c = wkt_c.split(", ")
+
+    # remove duplicate point
+    num_p = len(wkt_c) - 1
+    remove_index = []
+    for dup in sorted(list_duplicates(wkt_c)):
+        dup_index = dup[1]
+        if 0 in dup_index and num_p in dup_index and len(dup_index) == 2:
+            pass
+            
+        elif 0 in dup_index and num_p in dup_index and len(dup_index) > 2:
+            dup_index_num = len(dup_index)-1
+            for j in range(1, dup_index_num):
+                ri = dup_index[j]
+                remove_hindex.append(ri)
+                
+        else:
+            dup_index_num = len(dup_index)-1
+            for j in range(dup_index_num):
+                ri = dup_index[j]
+                remove_hindex.append(ri)
+
+    wall_f = []
+    wkt_f = []
+    for p in range(len(wkt_c)):
+        if p not in remove_index:
+            wkt_u = wkt_c[p]
+            wkt_f.append(wkt_u)
+
+            if p < (len(wkt_c)-1):
+                wall_u = wall_c[p]
+                wall_f.append(wall_u)
+
+    wkt_f = ", ".join(wkt_f)
+    wkt_f = "POLYGON ((" + wkt_f + "))"
+    
+    return wall_f, wkt_f
+
+
+
+def clean_geometry(wall, wkt):
+    # load geometry
+    geo = shapely.wkt.loads(wkt)
+    
+    # move to (0,0)
+    geo_centroid = geo.centroid
+    translation_vector = (-geo_centroid.x, -geo_centroid.y)
+    moved_coords = [(x + translation_vector[0], y + translation_vector[1]) for x, y in geo.exterior.coords]
+    moved_geo = shapely.wkt.loads('POLYGON ((' + ', '.join([f'{x} {y}' for x, y in moved_coords]) + '))')
+
+    # if counterclockwise
+    if moved_geo.exterior.is_ccw:
+        geo_ccw = moved_geo
+        wall_ccw = wall
+    else:
+        geo_ccw = shapely.geometry.polygon.orient(moved_geo, 1)
+        
+        walltypes = len(list(set(wall)))
+        if walltypes == 1:
+            wall_ccw = wall
+        else:
+            wall_ccw = wall[::-1]
+
+
+    # ccw_geo 
+    coor_ccw = geo_ccw.exterior.coords
+    coor_ccw = list(coor_ccw)
+    coor_ccw = coor_ccw[:-1]
+        
+    coor_ccw_num = len(coor_ccw)
+    coor_ccw_xpy_lst = []
+    for i in range(coor_ccw_num):
+        coor_ccw_x = coor_ccw[i][0]
+        coor_ccw_y = coor_ccw[i][1]
+        coor_ccw_xpy = coor_ccw_x + coor_ccw_y
+        coor_ccw_xpy_lst.append(coor_ccw_xpy)
+        
+    coor_ccw_xpy_min_index = np.array(coor_ccw_xpy_lst).argmin()
+    coor_ccw_sort_index = []
+    for i in range(len(coor_ccw_xpy_lst)):
+        index_max = len(coor_ccw_xpy_lst) - 1 - coor_ccw_xpy_min_index
+        if i <= index_max:
+            sort_index = coor_ccw_xpy_min_index + i
+        else:
+            sort_index =  i - len(coor_ccw_xpy_lst) + coor_ccw_xpy_min_index
+        coor_ccw_sort_index.append(sort_index)
+        
+    
+    coor_sort_lst = []
+    wall_sort_lst = []
+    for i in range(len(coor_ccw_sort_index)):
+        sort_index = coor_ccw_sort_index[i]
+        sort_coor = coor_ccw[sort_index]
+        sort_wall = wall_ccw[sort_index]
+        coor_sort_lst.append(sort_coor)
+        wall_sort_lst.append(sort_wall)
+        
+    geo_s = Polygon(coor_sort_lst)
+    wall_s = wall_sort_lst
+    return wall_s, geo_s
+
+    
+def segments(polyline):
+    return list(map(LineString, zip(polyline.coords[:-1], polyline.coords[1:])))
+
+
+
+
+def points4cv(x, y, xmin_abs, ymin_abs, scale):
+    points = []
+    for j in range(len(x)):
+        xp =x[j]
+        yp =y[j]
+
+        xp = (xp + xmin_abs +1) * scale
+        yp = (yp + ymin_abs +1) * scale
+        p = [int(xp), int(yp)]
+        points.append(p)
+
+    p_4_cv = np.array(points)
+    return p_4_cv
+
+
+
+
+def gridpoints(apa_geo, size):
+    latmin, lonmin, latmax, lonmax = apa_geo.bounds
+    prep_moved_apa_geo = prep(apa_geo)
+    
+    # construct a rectangular mesh
+    gp = []
+    for lat in np.arange(latmin, latmax, size):
+        for lon in np.arange(lonmin, lonmax, size):
+            gp.append(Point((round(lat,5), round(lon,5))))
+    gps = prep_moved_apa_geo.contains(gp)   
+    gpf = [i for indx,i in enumerate(gp) if gps[indx] == True]
+    grid_points = MultiPoint(gpf)
+    return grid_points
+
+
+
+def exterior_wall(apa_line, apa_wall):
+    apa_wall_O = [i for indx,i in enumerate(segments(apa_line)) if apa_wall[indx] == "O"]
+    apa_wall_O = MultiLineString(apa_wall_O)
+    return apa_wall_O
+
+
+def geo_coor(apa_geo):
+    apa_coor = apa_geo.exterior.coords
+    apa_coor = list(apa_coor)
+    return apa_coor
+
+
+
+

utils/edge_features.py

@@ -0,0 +1,113 @@
+import numpy as np
+import math
+
+import shapely
+from shapely.geometry import LineString
+
+import torch
+from torch_geometric.utils import to_undirected, remove_self_loops
+from torch_geometric.data import Data
+
+def segments(polyline):
+    return list(map(LineString, zip(polyline.coords[:-1], polyline.coords[1:])))
+
+
+def edge_graph(apa_line, apa_wall):
+    wall_seg = segments(apa_line)
+    num_seg = len(wall_seg)
+
+
+    edge_lst = []
+    for l in range(num_seg):
+        seg = wall_seg[l]
+        seg_length = seg.length
+        seg_pro = apa_wall[l]
+        south_cos = wall_segment_cosine("south", seg)
+        east_cos = wall_segment_cosine("east", seg)
+        north_cos = wall_segment_cosine("north", seg)
+        west_cos = wall_segment_cosine("west", seg)
+
+        if south_cos < 0:
+            south_cos = 0
+        if east_cos < 0:
+            east_cos = 0
+        if north_cos < 0:
+            north_cos = 0
+        if west_cos < 0:
+            west_cos = 0
+
+        if seg_pro == "I":
+            south_cos = 0
+            east_cos = 0
+            north_cos = 0
+            west_cos = 0
+            
+        if seg_pro == "O":
+            seg_boo = 1
+        else:
+            seg_boo = 0
+
+            
+
+        edge = [seg_boo, seg_length, south_cos, north_cos, west_cos, east_cos]
+        edge_lst.append(edge)
+
+
+    ms_lst = []
+    me_lst = []
+    for k in range(num_seg):
+        if k == (num_seg - 1):
+            ms = k
+            me = 0
+        else:
+            ms = k
+            me = k+1
+        ms_lst.append(ms)
+        me_lst.append(me)
+    mse = [ms_lst, me_lst]
+
+    datasets = []
+    for i in range(2):
+        node_features = torch.FloatTensor(edge_lst)
+        x = node_features
+
+        edge_index = torch.tensor(mse, dtype=torch.long)
+        edge_index, _ = remove_self_loops(edge_index)
+        edge_index = to_undirected(edge_index=edge_index)
+
+        data = Data(x=x, edge_index=edge_index)
+        datasets.append(data)
+    return datasets
+
+
+
+def wall_segment_cosine(direction, apa_line_seg):
+    seg_s = list(apa_line_seg.coords)[0]
+    seg_e = list(apa_line_seg.coords)[1]
+    
+    normal_x = seg_e[0] - seg_s[0]
+    normal_y = seg_e[1] - seg_s[1]
+    
+    normal_s = (-normal_y, normal_x)
+    normal_e = (normal_y, -normal_x)
+    
+    o = np.array([-normal_y, normal_x])
+    w = np.array([normal_y, -normal_x])
+    
+    if direction == "south":
+        d = np.array([-normal_y, normal_x-1])
+    if direction == "east":
+        d = np.array([-normal_y+1, normal_x])
+    if direction == "north":
+        d = np.array([-normal_y, normal_x+1])
+    if direction == "west":
+        d = np.array([-normal_y-1, normal_x])
+        
+    od = d - o
+    ow = w - o
+    
+    cosine = np.dot(od, ow) / (np.linalg.norm(od) * np.linalg.norm(ow))
+    return cosine
+
+
+    

utils/emb_model/Edge_64.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:127b70aeaa5ebe3eadc2bdf26f0521bc87cd78db3454087c86f8b12eecc9980c
+size 174179

utils/emb_model/Node_64.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c22fee85b8744349a7746eb61324f5259032fac224edb536aea06542a12cdd67
+size 173681

utils/node_features.py

@@ -0,0 +1,123 @@
+import numpy as np
+import math
+
+import shapely
+from shapely.geometry import Polygon, LineString
+
+import torch
+from torch_geometric.utils import to_undirected, remove_self_loops
+from torch_geometric.data import Data
+
+
+def node_graph(apa_coor, apa_geo):
+    num_op = len(apa_coor)
+    apa_coor = apa_coor[0:-1]
+    # apa_coor.pop(num_op-1)
+
+    node_lst = []
+    num_p = len(apa_coor)
+    for j in range(num_p):
+        p = apa_coor[j]
+        if j == 0:
+            sindex = -1
+            oindex = j
+            eindex = 1
+        elif j == (len(apa_coor)-1):
+            sindex = j-1
+            oindex = j
+            eindex = 0
+        else:
+            sindex = j-1
+            oindex = j
+            eindex = j+1
+
+        sp = apa_coor[sindex]
+        s = np.array(sp)
+
+        op = apa_coor[oindex]
+        o = np.array(op)
+        ox = op[0]
+        oy = op[1]
+
+        ep = apa_coor[eindex]
+        e = np.array(ep)
+
+        Area = apa_geo.area
+        local_polygon = Polygon((sp, op, ep))
+        larea = (local_polygon.area) / Area
+
+        se = LineString((sp, ep))
+        llength = se.length / math.sqrt(Area)
+
+        osv = s - o
+        oev = e - o
+
+        langle = angle_between(osv, oev)
+        if langle < 0:
+            langle = langle + (2*math.pi)
+
+
+        oop = (0, 0)
+        oo = np.array(oop)
+        regional_polygon = Polygon((sp, oop, ep))
+        regional_polygon_area = regional_polygon.area
+        rarea = regional_polygon_area / Area
+
+        regional_polygon_perimeter = regional_polygon.length / 2
+        rperimeter = regional_polygon_perimeter / math.sqrt(Area)
+
+        rradius = (regional_polygon_area / regional_polygon_perimeter) / math.sqrt(Area)
+
+        oosv = s - oo
+        ooev = e - oo
+        rangle = angle_between(oosv, ooev)
+        if rangle < 0:
+            rangle = rangle + (2*math.pi)
+
+        #ox, oy,
+        nl = [larea, llength, langle, rarea, rperimeter, rradius, rangle]
+
+        node_lst.append(nl)
+
+
+    ms_lst = []
+    me_lst = []
+    for k in range(num_p):
+        if k == (num_p - 1):
+            ms = k
+            me = 0
+        else:
+            ms = k
+            me = k+1
+        ms_lst.append(ms)
+        me_lst.append(me)
+    mse = [ms_lst, me_lst]
+
+    datasets = []
+    for i in range(2):
+        node_f = torch.FloatTensor(node_lst)
+        x = node_f
+
+        edge_index = torch.tensor(mse, dtype=torch.long)
+        edge_index, _ = remove_self_loops(edge_index)
+        edge_index = to_undirected(edge_index=edge_index)
+
+        data = Data(x=x, edge_index=edge_index)
+        datasets.append(data)
+    return datasets
+
+
+
+
+def angle_between(v1, v2):
+    """ Returns the angle in radians between vectors 'v1' and 'v2'
+        The sign of the angle is dependent on the order of v1 and v2
+        so acos(norm(dot(v1, v2))) does not work and atan2 has to be used, see:
+        https://stackoverflow.com/questions/21483999/using-atan2-to-find-angle-between-two-vectors
+    """
+    arg1 = np.cross(v1, v2)
+    arg2 = np.dot(v1, v2)
+    angle = np.arctan2(arg1, arg2)
+    return angle
+
+    

utils/shape_features.py

@@ -0,0 +1,625 @@
+import pandas as pd
+import numpy as np
+import math
+from statistics import mean, stdev
+from collections import defaultdict
+
+import shapely
+import shapely.wkt
+from shapely.geometry import Point, MultiPoint, LineString, MultiLineString, Polygon, LinearRing
+from shapely.ops import voronoi_diagram, substring, unary_union, nearest_points
+from shapely import affinity
+from shapely.prepared import prep
+
+import cv2 as cv
+
+
+def segments(polyline):
+    return list(map(LineString, zip(polyline.coords[:-1], polyline.coords[1:])))
+
+
+def scale_move_x(x, xmin_abs, scale):
+    xn = (x / scale) - 1 - xmin_abs
+    return xn
+
+def scale_move_y(y, ymin_abs, scale):
+    yn = (y / scale) - 1 - ymin_abs
+    return yn
+
+def scale_area(a, scale):
+    a = a / (scale**2)
+    return a
+
+def scale_perimeter(p, scale):
+    p = p / scale
+    return p
+
+
+def wall_segment_cosine(direction, apa_line_seg):
+    seg_s = list(apa_line_seg.coords)[0]
+    seg_e = list(apa_line_seg.coords)[1]
+    
+    normal_x = seg_e[0] - seg_s[0]
+    normal_y = seg_e[1] - seg_s[1]
+    
+    normal_s = (-normal_y, normal_x)
+    normal_e = (normal_y, -normal_x)
+    
+    o = np.array([-normal_y, normal_x])
+    w = np.array([normal_y, -normal_x])
+    
+    if direction == "south":
+        d = np.array([-normal_y, normal_x-1])
+    if direction == "east":
+        d = np.array([-normal_y+1, normal_x])
+    if direction == "north":
+        d = np.array([-normal_y, normal_x+1])
+    if direction == "west":
+        d = np.array([-normal_y-1, normal_x])
+        
+    od = d - o
+    ow = w - o
+    
+    cosine = np.dot(od, ow) / (np.linalg.norm(od) * np.linalg.norm(ow))
+    return cosine
+
+
+# Dir_S_longestedge, Dir_N_longestedge, Dir_W_longestedge, Dir_E_longestedge, Dir_S_max, Dir_N_max, Dir_W_max, Dir_E_max, Facade_length, Facade_ratio
+def wall_direction_ratio(apa_line, apa_wall):
+    apa_wall_O = [i for indx,i in enumerate(segments(apa_line)) if apa_wall[indx] == "O"]
+    apa_wall_O = MultiLineString(apa_wall_O)
+
+    wall_O_length = []
+    wall_O_south = []
+    wall_O_east = []
+    wall_O_north = []
+    wall_O_west = []
+    apa_wall_O_num = len(apa_wall_O.geoms)
+
+    if apa_wall_O_num > 0:
+        for i in range(apa_wall_O_num):
+            wall_seg = apa_wall_O.geoms[i]
+            wall_length = wall_seg.length
+            south_cos = wall_segment_cosine("south", wall_seg)
+            east_cos = wall_segment_cosine("east", wall_seg)
+            north_cos = wall_segment_cosine("north", wall_seg)
+            west_cos = wall_segment_cosine("west", wall_seg)
+
+            if south_cos < 0:
+                south_cos = 0
+            if east_cos < 0:
+                east_cos = 0
+            if north_cos < 0:
+                north_cos = 0
+            if west_cos < 0:
+                west_cos = 0
+
+            wall_O_length.append(wall_length)
+            wall_O_south.append(south_cos)
+            wall_O_east.append(east_cos)
+            wall_O_north.append(north_cos)
+            wall_O_west.append(west_cos)
+
+
+        max_length_index = np.array(wall_O_length).argmax()
+        Dir_S_longestedge =  wall_O_south[max_length_index]   
+        Dir_N_longestedge =  wall_O_north[max_length_index]   
+        Dir_W_longestedge =  wall_O_west[max_length_index]  
+        Dir_E_longestedge =  wall_O_east[max_length_index]
+
+        Dir_S_max = max(wall_O_south)
+        Dir_N_max = max(wall_O_north)
+        Dir_W_max = max(wall_O_west)
+        Dir_E_max = max(wall_O_east)
+
+        Facade_length = apa_wall_O.length
+        apa_line_length = apa_line.length
+        Facade_ratio = Facade_length / apa_line_length
+    else:
+        Dir_S_longestedge = 0
+        Dir_N_longestedge = 0
+        Dir_W_longestedge = 0
+        Dir_E_longestedge = 0
+        Dir_S_max = 0
+        Dir_N_max = 0
+        Dir_W_max = 0
+        Dir_E_max = 0
+        Facade_length = 0
+        Facade_ratio = 0
+
+    return Dir_S_longestedge, Dir_N_longestedge, Dir_W_longestedge, Dir_E_longestedge, Dir_S_max, Dir_N_max, Dir_W_max, Dir_E_max, Facade_length, Facade_ratio
+
+
+# apa_geo
+def apartment_perimeter(apa_geo):
+    perimeter =apa_geo.length
+    return perimeter
+
+
+
+def apartment_area(apa_geo):
+    area =apa_geo.area
+    return area
+
+
+def boundingbox(apa_geo):
+    boundingbox = apa_geo.bounds
+    return boundingbox
+
+
+# BBox_width_x, BBox_height_y, Aspect_ratio, Extent, ULC_x, ULC_y, LRC_x, LRC_y
+def boundingbox_features(apa_geo):
+    # [Aspect_ratio, Extent] ---> https://docs.opencv.org/3.4/d1/d32/tutorial_py_contour_properties.html
+
+    bbox_xy = boundingbox(apa_geo)
+    bbox_geo = Polygon([(bbox_xy[0], bbox_xy[1]), (bbox_xy[2], bbox_xy[1]), (bbox_xy[2], bbox_xy[3]), (bbox_xy[0], bbox_xy[3])])
+    
+    BBox_width_x = bbox_xy[2] - bbox_xy[0]
+    BBox_height_y = bbox_xy[3] - bbox_xy[1]
+    Aspect_ratio = BBox_width_x / BBox_height_y
+    
+    bbox_geo_area = bbox_geo.area
+    Area = apartment_area(apa_geo)
+    Extent = Area / bbox_geo_area
+    
+    ULC_x = bbox_xy[0]
+    ULC_y = bbox_xy[3]
+    LRC_x = bbox_xy[2]
+    LRC_y = bbox_xy[1]
+    
+    return BBox_width_x, BBox_height_y, Aspect_ratio, Extent, ULC_x, ULC_y, LRC_x, LRC_y
+
+
+# Max_diameter
+def max_diameter(apa_geo):
+    # [Max_diameter] ---> https://www.mvtec.com/doc/halcon/12/en/diameter_region.html
+    apa_coor = list(apa_geo.exterior.coords)
+    
+    pp_dis_lst = []
+    for i in apa_coor:
+        for j in apa_coor:
+            pp_dis = Point(i).distance(Point(j))
+            pp_dis_lst.append(pp_dis)
+
+    max_diameter = max(pp_dis_lst)
+    return max_diameter
+
+
+
+def fractality(apa_geo):
+    # [Fractality] ---> https://onlinelibrary.wiley.com/doi/epdf/10.1111/j.1538-4632.2000.tb00419.x
+    # Basaraner, M. and Cetinkaya, S. (2017) ‘Performance of shape indices and classification schemes for characterising perceptual shape complexity of building footprints in GIS’, International Journal of Geographical Information Science, 31(10), pp. 1952–1977. doi:10.1080/13658816.2017.1346257. 
+    Area = apartment_area(apa_geo)
+    Perimeter = apartment_perimeter(apa_geo)
+
+    fractality = 1 - ((math.log(Area) / (2 * math.log(Perimeter))))
+    return fractality
+
+
+
+def circularity(apa_geo):
+    # [Circularity] ---> https://www.mvtec.com/doc/halcon/12/en/circularity.html
+    apa_coor = list(apa_geo.exterior.coords)
+    op_dis_lst = []
+    for i in apa_coor:
+        op_dis = Point((0, 0)).distance(Point(i))
+        op_dis_lst.append(op_dis)
+        
+    Max_radius = max(op_dis_lst)
+    
+    Area = apartment_area(apa_geo)
+    
+    circularity = Area / ((math.pi) * (Max_radius**2))
+    return circularity
+
+
+
+def outer_radius(p_4_cv, xmin_abs, ymin_abs, scale):
+    # [Outer_radius] ---> https://docs.opencv.org/4.x/d3/dc0/group__imgproc__shape.html#ga8ce13c24081bbc7151e9326f412190f1
+    (xmin,ymin),radius = cv.minEnclosingCircle(p_4_cv)
+    mini_Enclosing_Cir_x = scale_move_x(xmin, xmin_abs, scale)
+    mini_Enclosing_Cir_y = scale_move_y(ymin, ymin_abs, scale)
+
+    mini_Enclosing_Cir_radius = scale_perimeter(radius, scale)
+    outer_radius = mini_Enclosing_Cir_radius
+    return outer_radius
+
+
+
+def inner_radius(apa_geo, apa_line):
+    # [Inner_radius] ---> https://www.sthu.org/blog/14-skeleton-offset-topology/index.html
+    dis_p = []
+    for i in np.arange(0, apa_line.length, 0.1):
+        s = substring(apa_line, i, i+0.1)
+        dis_p.append(s.boundary.geoms[0])
+    mp = MultiPoint(dis_p)
+
+    regions = voronoi_diagram(mp)
+
+    vo_p = []
+    for i in range(len(regions.geoms)):
+        vo = regions.geoms[i]
+        b = list(vo.exterior.coords)
+        for j in range(len(b)):
+            p = Point(b[j])
+            vo_p.append(p)
+    vo_p = MultiPoint(vo_p)
+    vo_p = unary_union(vo_p)
+    vo_p_b = []
+    for i in range(len(vo_p.geoms)):
+        t_c_p = vo_p.geoms[i]
+        pc = apa_geo.contains(t_c_p)
+        vo_p_b.append(pc)
+    vo_filtered_p = [i for indx,i in enumerate(vo_p.geoms) if vo_p_b[indx] == True]
+
+    vo_d = []
+    for i in range(len(vo_filtered_p)):
+        c = Point(vo_filtered_p[i])
+        d_min = c.distance(apa_line)
+        vo_d.append(d_min)
+
+    vo_r_max = max(vo_d)
+    vo_r_max_index = vo_d.index(vo_r_max)
+    vo_c_max = vo_filtered_p[vo_r_max_index]
+    vo_c_max = list(vo_c_max.coords)
+
+    max_Inner_Circle_x = vo_c_max[0][0]
+    max_Inner_Circle_y = vo_c_max[0][1]
+    max_Inner_Circle_r = vo_r_max
+    inner_radius = max_Inner_Circle_r
+    return inner_radius
+
+
+
+def roundness_features(apa_line):
+    # [Dist_mean, Dist_sigma, Roundness] ---> https://www.mvtec.com/doc/halcon/12/en/roundness.html
+    rou_p = []
+    for i in np.arange(0, apa_line.length, 0.5):
+        s = substring(apa_line, i, i+0.5)
+        rou_p.append(s.boundary.geoms[0])
+    rp = MultiPoint(rou_p)
+    
+    ro_dis = []
+    for i in range(len(rp.geoms)):
+        rpp = rp.geoms[i]
+        ro = Point(rpp).distance(Point((0, 0)))
+        ro_dis.append(ro)
+        
+    dist_mean = mean(ro_dis) 
+#     dist_sigma = stdev(ro_dis) 
+         
+    dev_lst = []
+    for i in ro_dis:
+        dev = (i - dist_mean)**2
+        dev_lst.append(dev)
+    dist_sigma = mean(dev_lst)
+    dist_sigma = math.sqrt(dist_sigma)
+    roundness = 1 - (dist_sigma/dist_mean)
+    
+    return dist_mean, dist_sigma, roundness
+
+
+def compactness(apa_geo):
+    # [Compactness] ---> https://fisherzachary.github.io/public/r-output.html
+    Area = apartment_area(apa_geo)
+    Perimeter = apartment_perimeter(apa_geo)
+    
+    compactness = (4*(math.pi)) * (Area / (Perimeter**2))
+    return compactness
+
+
+
+def equivalent_diameter(apa_geo):
+    # https://docs.opencv.org/4.x/d1/d32/tutorial_py_contour_properties.html
+    Area = apartment_area(apa_geo)
+    
+    equivalent_diameter = math.sqrt((4 * Area) / math.pi)
+    return equivalent_diameter
+
+
+
+
+def shape_membership_index(apa_line):
+    # [Shape_membership_index] ---> Basaraner, M. and Cetinkaya, S. (2017) ‘Performance of shape indices and classification schemes for characterising perceptual shape complexity of building footprints in GIS’, International Journal of Geographical Information Science, 31(10), pp. 1952–1977. doi:10.1080/13658816.2017.1346257. 
+
+    line_smi = LineString([(0, 0), (30, 0)])
+
+    numl = 30
+    line_rot_degree = 360 / numl
+    line_rot = []
+    for an in range(numl):
+        ang = an*line_rot_degree
+        lr = affinity.rotate(line_smi, ang, (0, 0))
+        line_rot.append(lr)
+    line_rot = MultiLineString(line_rot)
+    smip = shapely.intersection(apa_line, line_rot)
+
+
+    simo_dis = []
+    for i in range(len(smip.geoms)):
+        sim_p = smip.geoms[i]
+        simo = Point(sim_p).distance(Point((0, 0)))
+        simo_dis.append(simo)
+    sim_r_max = max(simo_dis)
+
+    simo_maxd = []
+    for j in simo_dis:
+        rmax_d = j / sim_r_max
+        simo_maxd.append(rmax_d)
+
+    simo_maxd_mean = mean(simo_maxd) 
+
+    simo_rad = []
+    for j in range(len(simo_dis)):
+        s = simo_dis[j]    
+
+        if j == (len(simo_dis) - 1):
+            nu = 0
+        else:
+            nu = j+1
+        e = simo_dis[nu]
+
+        if s <= e:
+            a = np.array([1,s])
+            b = np.array([0,s])
+            c = np.array([1,e])
+        else:
+            a = np.array([1,e])
+            b = np.array([0,e])
+            c = np.array([1,s])
+
+        ba = a - b
+        bc = c - b
+
+        cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
+        angle_rad = np.arccos(cosine_angle)
+
+        simo_rad.append(angle_rad)
+
+    simo_rad_min = min(simo_rad)
+    simo_rad_max = max(simo_rad)
+    simo_cos = math.cos(simo_rad_max - simo_rad_min)
+    shape_membership_index = simo_cos * simo_maxd_mean
+    return shape_membership_index
+
+
+def convexity(p_4_cv, apa_geo, xmin_abs, ymin_abs, scale):
+    # [Convexity] ---> Basaraner, M. and Cetinkaya, S. (2017) ‘Performance of shape indices and classification schemes for characterising perceptual shape complexity of building footprints in GIS’, International Journal of Geographical Information Science, 31(10), pp. 1952–1977. doi:10.1080/13658816.2017.1346257. 
+
+    hull = cv.convexHull(p_4_cv)
+    hull_x = []
+    hull_y = []
+    for h in range(len(hull)):
+        h_x = hull[h][0][0]
+        h_x = scale_move_x(h_x, xmin_abs, scale)
+        hull_x.append(h_x)
+
+        h_y = hull[h][0][1]
+        h_y = scale_move_y(h_y, ymin_abs, scale)
+        hull_y.append(h_y)
+
+    hull_xy = []
+    for i in range(len(hull_x)):
+        hx = hull_x[i]
+        hy = hull_y[i]
+        hull_xy.append((hx, hy))
+    hull_geo = Polygon(hull_xy)
+    Hull_area = hull_geo.area
+    
+    Area = apartment_area(apa_geo)
+    convexity = Area / Hull_area
+    return convexity, hull_geo
+
+
+
+def rectangle_features(p_4_cv, apa_geo, xmin_abs, ymin_abs, scale):
+    # [Rectangularity] ---> Basaraner, M. and Cetinkaya, S. (2017) ‘Performance of shape indices and classification schemes for characterising perceptual shape complexity of building footprints in GIS’, International Journal of Geographical Information Science, 31(10), pp. 1952–1977. doi:10.1080/13658816.2017.1346257. 
+    rect = cv.minAreaRect(p_4_cv)
+    miniRect_rotation_angle = rect[2]
+    box = cv.boxPoints(rect)
+    box = np.intp(box)
+
+    miniRect_x = []
+    miniRect_y = []
+    for b in range(len(box)):
+
+        b_x = box[b][0]
+        b_x = scale_move_x(b_x, xmin_abs, scale)
+        miniRect_x.append(b_x)
+
+        b_y = box[b][1]
+        b_y = scale_move_y(b_y, ymin_abs, scale)
+        miniRect_y.append(b_y)
+
+    miniRec_xy = []
+    for i in range(len(miniRect_x)):
+        minirecx = miniRect_x[i]
+        minirecy = miniRect_y[i]
+        miniRec_xy.append((minirecx, minirecy))
+    miniRect_geo = Polygon(miniRec_xy)
+    miniRect_area = miniRect_geo.area
+    
+    Area = apartment_area(apa_geo)
+    rectangularity = Area / miniRect_area
+    rect_phi = (miniRect_rotation_angle * math.pi) / 180
+
+    miniRect_line = miniRect_geo.boundary
+    miniRect_segments = segments(miniRect_line)
+
+    seg_len = []
+    for s in miniRect_segments:
+        seg_len.append(s.length)
+    rect_width = max(seg_len)
+    rect_height = min(seg_len)
+    return rectangularity, rect_phi, rect_width, rect_height
+
+
+def squareness(apa_geo):
+    # [Squareness] ---> Basaraner, M. and Cetinkaya, S. (2017) ‘Performance of shape indices and classification schemes for characterising perceptual shape complexity of building footprints in GIS’, International Journal of Geographical Information Science, 31(10), pp. 1952–1977. doi:10.1080/13658816.2017.1346257. 
+
+    Area = apartment_area(apa_geo)
+    Perimeter = apartment_perimeter(apa_geo)
+    
+    squareness = (4*(math.sqrt(Area))) / Perimeter
+    return squareness
+
+
+
+def moments(apa_geo):
+    # https://leancrew.com/all-this/2018/01/python-module-for-section-properties/
+    pts = list(apa_geo.exterior.coords)
+
+    if pts[0] != pts[-1]:
+        pts = pts + pts[:1]
+    x = [ c[0] for c in pts ]
+    y = [ c[1] for c in pts ]
+    sxx = syy = sxy = 0
+    a = apartment_area(apa_geo)
+    cx = apa_geo.centroid.x
+    cy = apa_geo.centroid.y
+    for i in range(len(pts) - 1):
+        sxx += (y[i]**2 + y[i]*y[i+1] + y[i+1]**2)*(x[i]*y[i+1] - x[i+1]*y[i])
+        syy += (x[i]**2 + x[i]*x[i+1] + x[i+1]**2)*(x[i]*y[i+1] - x[i+1]*y[i])
+        sxy += (x[i]*y[i+1] + 2*x[i]*y[i] + 2*x[i+1]*y[i+1] + x[i+1]*y[i])*(x[i]*y[i+1] - x[i+1]*y[i])
+    return sxx/12 - a*cy**2, syy/12 - a*cx**2, sxy/24 - a*cx*cy
+
+
+def moment_index(apa_geo, Convexity, Compactness):
+	# https://www.researchgate.net/publication/228557311_A_COMBINED_AUTOMATED_GENERALIZATION_MODEL_BASED_ON_THE_RELATIVE_FORCES_BETWEEN_SPATIAL_OBJECTS
+	Ixx, Iyy, Ixy = moments(apa_geo)
+	ratio = max(Ixx, Iyy) / min(Ixx, Iyy)
+	# Convexity, Hull_geo = convexity(p_4_cv)
+	# Compactness = compactness(apa_geo)
+	moment_index = (Convexity * Compactness) / ratio
+	return moment_index
+
+
+
+def ndetour_index(apa_geo, Hull_geo):
+    # [nDetour_index] ---> Basaraner, M. and Cetinkaya, S. (2017) ‘Performance of shape indices and classification schemes for characterising perceptual shape complexity of building footprints in GIS’, International Journal of Geographical Information Science, 31(10), pp. 1952–1977. doi:10.1080/13658816.2017.1346257. 
+
+    Hull_line = Hull_geo.boundary
+    Hull_length = Hull_line.length
+    Area = apartment_area(apa_geo)
+    ndetour_index = (2 * math.sqrt(Area * math.pi)) / Hull_length
+    return ndetour_index
+
+
+def ncohesion_index(apa_geo, grid_points):
+    # [nCohesion_index] ---> Basaraner, M. and Cetinkaya, S. (2017) ‘Performance of shape indices and classification schemes for characterising perceptual shape complexity of building footprints in GIS’, International Journal of Geographical Information Science, 31(10), pp. 1952–1977. doi:10.1080/13658816.2017.1346257. 
+
+    grid_p = grid_points.geoms
+    grid_n = len(grid_p)
+    gg_dis_lst = []
+    for i in grid_p:
+        for j in grid_p:
+            gg_dis = Point(i).distance(Point(j))
+            gg_dis_lst.append(gg_dis)
+            
+    Area = apartment_area(apa_geo)
+    ncohesion_index = (0.9054 * math.sqrt(Area / math.pi)) / (sum(gg_dis_lst) / (grid_n * (grid_n-1)))
+    return ncohesion_index
+
+
+
+def nproximity_nspin_index(apa_geo, grid_points):
+    grid_p = grid_points.geoms
+    
+    go_dis_lst = []
+    for i in grid_p:
+        go_dis = Point(i).distance(Point(0,0))
+        go_dis_lst.append(go_dis)
+
+    go_dis_mean = mean(go_dis_lst)
+    Area = apartment_area(apa_geo)
+    nproximity_index = ((2 / 3) * math.sqrt(Area / math.pi)) / go_dis_mean
+    
+    nspin_index = (0.5 * (Area / math.pi)) / (go_dis_mean**2)
+    
+    return nproximity_index, nspin_index
+
+
+
+def nexchange_index(apa_geo):
+    Area = apartment_area(apa_geo)
+
+    eac_r = math.sqrt(Area / math.pi) 
+    eac = Point(0,0).buffer(eac_r)
+    eac_inter = apa_geo.intersection(eac)
+
+    if eac_inter.geom_type == "Polygon":
+        eac_area = eac_inter.area
+    else:
+        eacga_lst = []
+        for i in range(len(eac_inter.geoms)):
+            eacg = eac_inter.geoms[i]
+            eacga = eacg.area
+            eacga_lst.append(eacga)
+        eac_area = sum(eacga_lst)
+    nexchange_index = eac_area / Area
+    return nexchange_index
+
+
+
+def nperimeter_index(apa_geo):
+    Area = apartment_area(apa_geo)
+    Perimeter = apartment_perimeter(apa_geo)
+
+    nperimeter_index = (2 * math.sqrt(math.pi * Area)) / Perimeter
+    return nperimeter_index
+
+
+
+def ndepth_index(apa_geo, apa_line, grid_points):
+    moved_apa_line = apa_line
+    grid_p = grid_points.geoms
+    
+    nea_len_lst = []
+    for i in grid_p:
+        nea_line = LineString(nearest_points(moved_apa_line, i))
+        nea_len = nea_line.length
+        nea_len_lst.append(nea_len)
+    nea_len_mean = mean(nea_len_lst)
+
+    Area = apartment_area(apa_geo)
+    ndepth_index = (3 * nea_len_mean) / math.sqrt(Area / math.pi)
+    return ndepth_index
+
+
+
+def ngirth_index(apa_geo, Inner_radius):
+    Area = apartment_area(apa_geo)
+
+    ngirth_index = Inner_radius / math.sqrt(Area / math.pi)
+    return ngirth_index
+
+
+
+def nrange_index(apa_geo, Outer_radius):
+    Area = apartment_area(apa_geo)
+
+    nrange_index =  math.sqrt(Area / math.pi) / Outer_radius
+    return nrange_index
+
+
+
+def ntraversal_index(apa_geo, apa_line):
+    rou_p = []
+    for i in np.arange(0, apa_line.length, 0.5):
+        s = substring(apa_line, i, i+0.5)
+        rou_p.append(s.boundary.geoms[0])
+    rp = MultiPoint(rou_p)
+    
+    rp_n = len(rp.geoms)
+    bb_dis_lst = []
+    for i in rp.geoms:
+        for j in rp.geoms:
+            bb_dis = Point(i).distance(Point(j))
+            bb_dis_lst.append(bb_dis)
+    bb_dis_mean = sum(bb_dis_lst) / (rp_n * (rp_n-1))
+    
+    Area = apartment_area(apa_geo)
+    
+    ntraversal_index = (4 * (math.sqrt(Area / math.pi) / math.pi)) / bb_dis_mean
+    return ntraversal_index
+