Grid Distance Function

grid_dist.py

"""
# Grid Dist
what's a loss function that incorporates spatial distance?

e.g. i want...
loss(y..., y...) == 0
0 < loss(yt, y1) < loss(yt, y2)
loss(yt, y2) == loss(yt, y2.T)
Tweet: https://twitter.com/TheShubhanshu/status/1482132766607265796
"""

import numpy as np

class GridDist(object):
  def __init__(self, gd_dims=(3, 3), alpha=2):
    self.n_dims = len(gd_dims)
    x_size, y_size = gd_dims
    self.x = np.arange(-(x_size//2), x_size//2 + 1)[np.newaxis,:]
    self.y = np.arange(-(y_size//2), y_size//2 + 1)[:, np.newaxis]
    self.alpha = alpha
  
  def __call__(self, yt, y1):
    y_diff = yt - y1
    dist = (
      (self.x*y_diff).sum()**self.alpha 
    + (self.y*y_diff).sum()**self.alpha
    )
    return dist
  
class GeneralGridDist(object):
  def __init__(self, gd_dims=(3, 3), alpha=2):
    self.n_dims = len(gd_dims)
    self.dim_points = [
      np.arange(-(d//2), d//2 + 1).reshape([d if j == i else 1 for j in range(self.n_dims)])
      for i, d in enumerate(gd_dims)
    ]
    self.alpha = alpha
  
  def __call__(self, yt, y1):
    y_diff = yt - y1
    dist = np.sum([
      (d*y_diff).sum()**self.alpha
      for d in self.dim_points
    ])
    return dist
  
  
## Usage

yt = np.zeros((3, 3))
y1 = np.zeros((3, 3))
y2 = np.zeros((3, 3))
yt[0, 0] = 1
y1[0, 2] = 1
y2[1, 1] = 1

gd = GridDist()
gd(yt, y1), gd(yt, y2), gd(yt, yt), gd(yt, y2.T), gd(yt, y1.T)

## 3D grid

yt = np.zeros((3, 3, 3))
y1 = np.zeros((3, 3, 3))
y2 = np.zeros((3, 3, 3))
yt[0, 0, 0] = 1
y1[0, 0, 2] = 1
y2[1, 1, 1] = 1

ggd = GeneralGridDist(gd_dims=(3, 3, 3))
gd(yt, y1), gd(yt, y2), gd(yt, yt), gd(yt, y2.T), gd(yt, y1.T)