Introduction

TorchSpatial is a PyTorch package for Spatial Representation Learning. Currently, it provides 15 commonly recognized location encoders, including location encoders built on projected 2D space and 3D space. It is flexible enough to support the development of any location encoders. In the future, we plan also to support other spatial data types such as polylines, polygons, spatial networks, etc.

2D Location Encoder

For location encoders applicable in a projected 2D space, TorchSpatial implements the following 7 models:

  • tile is a discretized-based location encoder employed by many previous studies, which splits geographic areas into grids and uses grid embedding to represent locations within each grid. See details in tile.

  • wrap uses a coordinate wrap mechanism to convert each dimension of location into 2 numbers and feed them into \(\mathbf{NN}^{\mathit{wrap}}()\). \(\mathbf{NN}^{\mathit{wrap}}()\) consists of four residual blocks which are implemented as linear layers. See details in wrap.

  • wrap + ffn is a variant of wrap that substitutes \(\mathbf{NN}^{\mathit{wrap}}()\) in wrap with \(\mathbf{NN}^{\mathit{ffn}}()\). See details in wrap-ffn.

  • rbf randomly samples \(W\) points from the training dataset as RBF anchor points, and uses Gaussian kernels \(\exp\left(-\frac{\lVert x_i - x_{anchor}\rVert_2^2}{2\sigma^2}\right)\) on each anchor points, where \(\sigma\) is the kernel size. Each input point \(x_i\) is encoded as a \(W\)-dimension RBF feature vector, which is fed into \(\mathbf{NN}^{\mathit{ffn}}()\) to obtain the location embedding. See details in rbf.

  • rff means Random Fourier Features . It first encodes location \(x\) into a \(W\) dimension vector - \(\mathit{PE}^{\mathrm{rff}}(x) = \phi(x) = \frac{\sqrt{2}}{\sqrt{W}}\bigcup_{i=1}^W\left[\cos\left(\omega_i^T x + b_i\right)\right]\) where \(\omega_i \overset{\mathrm{i.i.d}}{\sim} \mathcal{N}(0, \delta^2 I)\) is a direction vector whose each dimension is independently sampled from a normal distribution. \(b_i\) is a shift value uniformly sampled from \([0, 2\pi]\) and \(I\) is an identity matrix. Each component of \(\phi(x)\) first projects \(x\) into a random direction \(w_i\) and makes a shift by \(b_i\). Then it wraps this line onto a unit circle in \(\mathbb{R}^2\) with the cosine function. \(\mathit{PE}^{\mathrm{rff}}(x)\) is further fed into \(\mathbf{NN}^{\mathit{ffn}}()\) to get a location embedding. See details in rff.

  • Space2Vec-grid and Space2Vec-theory are two multi-scale location encoder on 2D Euclidean space. Both of them implement the position encoder \(\mathit{PE}(x)\) as a deterministic Fourier mapping layer which is further fed into the \(\mathbf{NN}^{\mathit{ffn}}()\). Both models’ position encoders can be treated as performing a Fourier transformation on a 2D Euclidean space. See details in Space2Vec-grid and Space2Vec-theory.

3D Location Encoder

We also encompass 8 location encoders that learn location embeddings from 3D space as follows:

  • xyz first uses position encoder \(\mathit{PE}^{\mathrm{xyz}}(x)\) to convert the lat-lon spherical coordinates into 3D Cartesian coordinates centered at the sphere center. And then it feeds the 3D coordinates into a multilayer perceptron \(\mathbf{NN}^{\mathit{ffn}}()\). See details in xyz.

  • NeRF can be treated as a multiscale version of xyz using Neural Radiance Fields (NeRF) for its position encoder. See details in NeRF.

  • Sphere2Vec-sphereC, Sphere2Vec-sphereC+, Sphere2Vec-sphereM, Sphere2Vec-sphereM+, Sphere2Vec-dfs are variants of Sphere2Vec, a multi-scale location encoder for spherical surface. They are the first location encoder series that preserves the spherical surface distance between any two points to our knowledge. See details in Space2Vec-sphereC, Space2Vec-sphereC+, Space2Vec-sphereM, Space2Vec-sphereM+, and Space2Vec-dfs.

  • Siren(SH) is also a spherical location encoder proposed recently. It uses spherical harmonic basis functions as the position encoder \(\mathit{PE}^{\mathit{Spherical-Harmonics}}(x)\) and a sinusoidal representation network (SirenNets) as the \(\mathbf{NN}()\). See details in Siren(SH).