Utils

DropoutLayer

Bases: Module

Customized dropout layer where the nodes values are dropped with probability p.

Attributes:

Name	Type	Description
p	float	The drop probability

Source code in geospaNN/utils.py

class DropoutLayer(torch.nn.Module):
    """
    Customized dropout layer where the nodes values are dropped with probability p.

    Attributes:
        p (float):
            The drop probability
    """
    def __init__(self, p: float):
        super().__init__()
        self.p = p

    def forward(self, input):
        if self.training:
            u1 = (np.random.rand(*input.shape)<self.p) / self.p
            u1 *= u1
            return u1
        else:
            input *= self.p

EarlyStopping

Early stopping to stop the training when the loss does not improve after certain epochs.

Attributes:

Name	Type	Description
patience	int	How many epochs to wait before stopping when loss is not improving
min_delta	float	Minimum difference between new loss and old loss for new loss to be considered as an improvement

Source code in geospaNN/utils.py

class EarlyStopping():
    """
    Early stopping to stop the training when the loss does not improve after
    certain epochs.

    Attributes:
        patience (int):
            How many epochs to wait before stopping when loss is not improving
        min_delta (float):
            Minimum difference between new loss and old loss for new loss to be considered as an improvement
    """
    def __init__(self,
                 patience: Optional[int] = 5,
                 min_delta: Optional[float] = 0):
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = None
        self.early_stop = False
    def __call__(self, val_loss):
        if self.best_loss == None:
            self.best_loss = val_loss
        elif self.best_loss - val_loss > self.min_delta or math.isnan(self.best_loss):
            self.best_loss = val_loss
            # reset counter if validation loss improves
            self.counter = 0
        elif self.best_loss - val_loss < self.min_delta:
            self.counter += 1
            #print(f"INFO: Early stopping counter {self.counter} of {self.patience}")
            if self.counter >= self.patience:
                print('INFO: Early stopping')
                self.early_stop = True

LRScheduler

Learning rate scheduler. If the validation loss does not decrease for the given number of patience epochs, then the learning rate will decrease by by given factor.

Attributes:

Name	Type	Description
optimizer		the optimizer we are using
patience	int	How many epochs to wait before updating the lr. Default being 5.
min_lr	float	Least lr value to reduce to while updating.
factor	float	Factor by which the lr should be updated, i.e. new_lr = old_lr * factor.

Source code in geospaNN/utils.py

class LRScheduler():
    """
    Learning rate scheduler. If the validation loss does not decrease for the
    given number of `patience` epochs, then the learning rate will decrease by
    by given `factor`.

    Attributes:
        optimizer:
            the optimizer we are using
        patience (int):
            How many epochs to wait before updating the lr. Default being 5.
        min_lr (float):
            Least lr value to reduce to while updating.
        factor (float):
            Factor by which the lr should be updated, i.e. new_lr = old_lr * factor.
    """
    def __init__(
        self,
            optimizer,
            patience: Optional[int] = 5,
            min_lr: Optional[float] = 1e-6,
            factor: Optional[float] = 0.5
    ):
        self.optimizer = optimizer
        self.patience = patience
        self.min_lr = min_lr
        self.factor = factor
        self.lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                self.optimizer,
                mode='min',
                patience=self.patience,
                factor=self.factor,
                min_lr=self.min_lr,
                verbose=True
            )
    def __call__(self, val_loss):
        self.lr_scheduler.step(val_loss)

NNGP_cov

Bases: Sparse_B

A subclass of Sparse_B designed for the decorrelation using NNGP. The whole object is an NNGP approximation of the inverse square-root of a covariance matrix \Sigma. i.d. F^{-1/2}(I-B) ~ \Sigma^{-1/2} ...

Attributes:

Name	Type	Description
F_diag		torch.Tensor A vector of length n representing the diagonal matrix F.

Methods:

Name	Description
correlate	Approximately correlate the matrix X to \Sigma^{1/2}X.
decorrelate	Approximately decorrelate the matrix X to \Sigma^{-1/2}X.

Source code in geospaNN/utils.py

class NNGP_cov(Sparse_B):
    """
    A subclass of Sparse_B designed for the decorrelation using NNGP. The whole object is an NNGP approximation of the
    inverse square-root of a covariance matrix \Sigma. i.d. F^{-1/2}(I-B) ~ \Sigma^{-1/2}
    ...

    Attributes:
        F_diag : torch.Tensor
            A vector of length n representing the diagonal matrix F.

    Methods:
        correlate(x):
            Approximately correlate the matrix X to \Sigma^{1/2}X.

        decorrelate(x):
            Approximately decorrelate the matrix X to \Sigma^{-1/2}X.
    """
    def __init__(self, B, F_diag, Ind_list):
        super().__init__(B = B, Ind_list = Ind_list)
        assert len(F_diag) == B.shape[0]
        self.F_diag = F_diag

    def correlate(self, x: torch.Tensor
                  ) -> torch.Tensor:
        """
        Approximately correlate the matrix X to \Sigma^{1/2}X by calculating (I_B)^{-1}F^{1/2}X.

        Parameters:
            x : Input tensor X

        Returns:
            x_cor: Correlated X
        """
        assert x.shape[0] == self.n
        return self.invmul(torch.sqrt(self.F_diag) * x)

    def decorrelate(self, x: torch.Tensor
                    ) -> torch.Tensor:
        """
        Approximately decorrelate the matrix X to \Sigma^{-1/2}X by calculating F^{-1/2}(I_B)X.

        Parameters:
            x : Input tensor X

        Returns:
            x_decor: Decorrelated X
        """
        assert x.shape[0] == self.n
        return torch.sqrt(torch.reciprocal(self.F_diag))*self.matmul(x)

correlate(x)

Approximately correlate the matrix X to \Sigma^{1/2}X by calculating (I_B)^{-1}F^{1/2}X.

Parameters:

Name	Type	Description	Default
x		Input tensor X	required

Returns:

Name	Type	Description
x_cor	Tensor	Correlated X

Source code in geospaNN/utils.py

def correlate(self, x: torch.Tensor
              ) -> torch.Tensor:
    """
    Approximately correlate the matrix X to \Sigma^{1/2}X by calculating (I_B)^{-1}F^{1/2}X.

    Parameters:
        x : Input tensor X

    Returns:
        x_cor: Correlated X
    """
    assert x.shape[0] == self.n
    return self.invmul(torch.sqrt(self.F_diag) * x)

decorrelate(x)

Approximately decorrelate the matrix X to \Sigma^{-1/2}X by calculating F^{-1/2}(I_B)X.

Parameters:

Name	Type	Description	Default
x		Input tensor X	required

Returns:

Name	Type	Description
x_decor	Tensor	Decorrelated X

Source code in geospaNN/utils.py

def decorrelate(self, x: torch.Tensor
                ) -> torch.Tensor:
    """
    Approximately decorrelate the matrix X to \Sigma^{-1/2}X by calculating F^{-1/2}(I_B)X.

    Parameters:
        x : Input tensor X

    Returns:
        x_decor: Decorrelated X
    """
    assert x.shape[0] == self.n
    return torch.sqrt(torch.reciprocal(self.F_diag))*self.matmul(x)

Sparse_B

A sparse representation of the lower-triangular neighboring nxn matrix B_dense, where each row contains at most p non-zero values.

Attributes:

Name	Type	Description
B		nxp array contains all non-zero values in B_dense.
n		The number of samples.
neighbor_size		i.e. k in the documentation, the largest number of non-zero values in each row of B_dense.
Ind_list		The nxp index array indicating the location where values in B was in B_dense. For example, the [i,j]'s index is k means that B_dense[i,k] = B[i,j].

Methods:

Name	Description
to_numpy	Transform B to np.array
to_tensor	Transform B to torch.Tensor
matmul	Calculate the matrix product of B_dense[idx,:] and X
invmul	Calculate the matrix-vector product of B_dense^{-1} and y. Only used when the diagonal of B_dense is constantly 1.
Fmul	Return a new Sparse_B object where B_dense is replaced by the matrix product of F * B_dense. F_diag is the vector representation of the nxn diagonal matrix F.
to_dense	Return the dense form of B_dense as an np.array object.

Source code in geospaNN/utils.py

class Sparse_B():
    """
    A sparse representation of the lower-triangular neighboring nxn matrix B_dense,
    where each row contains at most p non-zero values.

    Attributes:
        B :
            nxp array contains all non-zero values in B_dense.
        n : 
            The number of samples.
        neighbor_size :
            i.e. k in the documentation, the largest number of non-zero values in each row of B_dense.
        Ind_list :
            The nxp index array indicating the location where values in B was in B_dense. For example, the [i,j]'s index is k
            means that B_dense[i,k] = B[i,j].

    Methods:
        to_numpy():
            Transform B to np.array

        to_tensor():
            Transform B to torch.Tensor

        matmul(X, idx = None):
            Calculate the matrix product of B_dense[idx,:] and X

        invmul(y):
            Calculate the matrix-vector product of B_dense^{-1} and y. Only used when the diagonal of B_dense is constantly 1.

        Fmul(F_diag):
            Return a new Sparse_B object where B_dense is replaced by the matrix product of F * B_dense.
            F_diag is the vector representation of the nxn diagonal matrix F.

        to_dense():
            Return the dense form of B_dense as an np.array object.
    """
    def __init__(self,
                 B: torch.Tensor | np.array,
                 Ind_list: np.array):
        self.B = B
        self.n = B.shape[0]
        self.neighbor_size = B.shape[1]
        self.Ind_list = Ind_list.astype(int)

    def to_numpy(self):
        if torch.is_tensor(self.B):
           self.B = self.B.detach().numpy()
        return self

    def to_tensor(self):
        if isinstance(self.B, np.ndarray):
            self.B = torch.from_numpy(self.B).float()
        return self

    def matmul(self, X, idx = None):
        if idx == None: idx = np.array(range(self.n))
        if torch.is_tensor(X):
            self.to_tensor()
            result = torch.empty((len(idx)))
            for k in range(len(idx)):
                i = idx[k]
                ind = self.Ind_list[i,:][self.Ind_list[i,:] >= 0]
                result[k] = torch.dot(self.B[i,range(len(ind))].squeeze(),X[ind])
        elif isinstance(X, np.ndarray):
            self.to_numpy()
            if np.ndim(X) == 1:
                result = np.empty((len(idx)))
                for k in range(len(idx)):
                    i = idx[k]
                    ind = self.Ind_list[i, :][self.Ind_list[i, :] >= 0]
                    result[k] = np.dot(self.B[i, range(len(ind))].reshape(-1), X[ind])
            elif np.ndim(X) == 2:
                result = np.empty((len(idx), X.shape[1]))
                for k in range(len(idx)):
                    i = idx[k]
                    ind = self.Ind_list[i, :][self.Ind_list[i, :] >= 0]
                    #result[i,:] = np.dot(self.B[i, range(len(ind))].reshape(-1), C_Ni[ind, :])
                    result[k,:] = np.dot(self.B[i, range(len(ind))].reshape(-1), X[ind,:])
        return result

    def invmul(self, y):
        if isinstance(y, np.ndarray):
            y = torch.from_numpy(y).float()
        y = y.reshape(-1)
        assert self.n == y.shape[0]
        x = y[0].unsqueeze(0)
        assert (self.B[:,0] == torch.ones(self.n)).all(), 'Only applies to I-B matrix'
        Indlist = self.Ind_list[:, 1:]
        B = -self.B[:, 1:]
        for i in range(1, self.n):
            ind = Indlist[i, :]
            id = ind >= 0
            if sum(id) == 0:
                x = torch.cat((x, y[i].unsqueeze(0)), dim = -1)
            else:
                x = torch.cat((x, (y[i] + torch.dot(x[ind[id]], B[i, :][id])).unsqueeze(0)), dim = -1)
        return x

    def Fmul(self, F_diag):
        res = Sparse_B(self.B.copy(), self.Ind_list.copy())
        for i in range(self.n):
            res.B[i,:] = F_diag[i]*self.B[i,:]
        return res

    def to_dense(self):
        B = np.zeros((self.n, self.n))
        for i in range(self.n):
            ind = self.Ind_list[i, :][self.Ind_list[i, :] >= 0]
            if len(ind) == 0:
                continue
            B[i, ind] = self.B[i,range(len(ind))]
        return B

Simulation(n, p, neighbor_size, fx, theta, coord=None, range=[0, 1], X_pattern='uniform', sparse=True)

Simulate spatial data

Simulate the spatial data based on the following model: Y(s) = f(X(s)) + w(s) + delta(s), where s are the spatial locations, X(s) is the spatial covariates, w(s) is the spatial effect (correlated noise), and delta(s) is the i.i.d random noise (nuggets).

Parameters:

Name	Type	Description	Default
n	int	Sample size.	required
p	int	Dimension of covariates.	required
fx	Callable	Function for the covariates' effect.	required
theta	tuple[float, float, float]	theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family, used for generating spatial random effect.	required
coord	Optional[tensor]	A nxd tensor as the locations where to simulate the data, if not specified, randomly sample from [range[0], range[1]]^2 square.	None
range	tuple[float, float]	A tuple [a, b] as the range of spatial locations. The spatial coordinates are sampled uniformly from [a, b]^2.	[0, 1]
sparse	Optional[bool]	To be implemented. Mainly interact with the rmvn() function.	True

Returns:

Name	Type	Description
X	Tensor	torch.Tensor nxp array sampled uniformly from [0,1] as the covariates.
Y	Tensor	torch.Tensor Length n vector as the observations consists of fixed covariate effect, spatial random effect, and random noise.
coord	Tensor	torch.Tensor Simulated spatial locations.
cov	Tensor | NNGP_cov	Covariance matrix based on the simulated coordinates.
corerr	Tensor	Simulated spatial random effect.

Source code in geospaNN/utils.py

def Simulation(n: int, p:int,
               neighbor_size: int,
               fx: Callable,
               theta: tuple[float, float, float],
               coord: Optional[torch.tensor] = None,
               range: tuple[float, float] = [0,1],
               X_pattern: Optional[str] = "uniform",
               sparse: Optional[bool] = True
               ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | NNGP_cov, torch.Tensor]:
    """Simulate spatial data

    Simulate the spatial data based on the following model: Y(s) = f(X(s)) + w(s) + delta(s),
    where s are the spatial locations, X(s) is the spatial covariates, w(s) is the spatial effect (correlated noise),
    and delta(s) is the i.i.d random noise (nuggets).

    Parameters:
        n:
            Sample size.
        p:
            Dimension of covariates.
        fx:
            Function for the covariates' effect.
        theta:
            theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family,
            used for generating spatial random effect.
        coord:
            A nxd tensor as the locations where to simulate the data, if not specified, randomly sample from [range[0], range[1]]^2
            square.
        range:
            A tuple [a, b] as the range of spatial locations. The spatial coordinates are sampled uniformly from [a, b]^2.
        sparse:
            To be implemented. Mainly interact with the rmvn() function.

    Returns:
        X: torch.Tensor
            nxp array sampled uniformly from [0,1] as the covariates.
        Y: torch.Tensor
            Length n vector as the observations consists of fixed covariate effect, spatial random effect, and random noise.
        coord: torch.Tensor
            Simulated spatial locations.
        cov:
            Covariance matrix based on the simulated coordinates.
        corerr:
            Simulated spatial random effect.
    """
    if coord is None:
        coord = (range[1] - range[0]) * torch.rand(n, 2) + range[0]
    sigma_sq, phi, tau = theta
    tau_sq = tau * sigma_sq

    cov = make_cov(coord, theta, neighbor_size)
    if X_pattern is "uniform":
        X = torch.rand(n, p)
    elif X_pattern is "correlated":
        _, _, _, _, X = Simulation(n*p, 1, neighbor_size, fx,
                                   torch.tensor([theta[0], 5*theta[1], theta[2]]), range=[0, 1])
        X = X.reshape(-1, p)
        X = (X - X.min()) / (X.max() - X.min())
    corerr = rmvn(torch.zeros(n), cov, sparse)
    Y = fx(X).reshape(-1) + corerr + torch.sqrt(tau_sq) * torch.randn(n)

    return X, Y, coord, cov, corerr

distance(coord1, coord2)

Distance matrix between two sets of points

Calculate the pairwise distance between two sets of locations.

Parameters:

Name	Type	Description	Default
coord1	Tensor	The nxd coordinates array for the first set.	required
coord12		The nxd coordinates array for the second set.	required

Returns:

Name	Type	Description
dist	Tensor	The distance matrix.

Source code in geospaNN/utils.py

def distance(coord1: torch.Tensor,
             coord2: torch.Tensor
             ) -> torch.Tensor:
    """Distance matrix between two sets of points

    Calculate the pairwise distance between two sets of locations.

    Parameters:
        coord1:
            The nxd coordinates array for the first set.
        coord12:
            The nxd coordinates array for the second set.

    Returns:
        dist:
            The distance matrix.
    """
    if coord1.ndim == 1:
        m = 1
        coord1 = coord1.unsqueeze(0)
    else:
        m = coord1.shape[0]
    if coord2.ndim == 1:
        n = 1
        coord2 = coord2.unsqueeze(0)
    else:
        n = coord2.shape[0]

    #### Can improve (resolved)
    coord1 = coord1.unsqueeze(0)
    coord2 = coord2.unsqueeze(1)
    dists = torch.sqrt(torch.sum((coord1 - coord2) ** 2, axis=-1))
    return dists

distance_np(coord1, coord2)

The numpy version of distance()

Calculate the pairwise distance between two sets of locations.

Parameters:

Name	Type	Description	Default
coord1	array	The nxd coordinates array for the first set.	required
coord2	array	The nxd coordinates array for the second set.	required

Returns:

Name	Type	Description
dists	array	The distance matrix.

See Also

distance : Distance matrix between two sets of points

Source code in geospaNN/utils.py

def distance_np(coord1: np.array,
                coord2: np.array
                ) -> np.array:
    """The numpy version of distance()

    Calculate the pairwise distance between two sets of locations.

    Parameters:
        coord1:
            The nxd coordinates array for the first set.
        coord2:
            The nxd coordinates array for the second set.

    Returns:
        dists:
            The distance matrix.

    See Also:
        distance : Distance matrix between two sets of points
    """
    m = coord1.shape[0]
    n = coord2.shape[0]
    #### Can improve (resolved)
    coord1 = coord1
    coord2 = coord2
    dists = np.zeros((m, n))
    for i in range(m):
        dists[i, :] = np.sqrt(np.sum((coord1[i] - coord2) ** 2, axis=1))
    return dists

krig_pred(w_train, coord_train, coord_test, theta, neighbor_size=20, q=0.95)

Kriging prediction (Gaussian process regression) with confidence interval.

Kriging prediction on testing locations based on the observations on the training locations. The kriging procedure assumes the observations are sampled from a Gaussian process, which is paramatrized here to have an exponential covariance structure using theta = [sigma^2, phi, tau]. NNGP appriximation is involved for efficient computation of matrix inverse. The conditional variance (kriging variance) is used to build the confidence interval using the quantiles (a/2, 1-a/2). (see https://arxiv.org/abs/2304.09157, section 4.3 for more details.)

Parameters:

Name	Type	Description	Default
w_train	Tensor	Training observations of the spatial random effect without any fixed effect.	required
coord_train	Tensor	Spatial coordinates of the training observations.	required
coord_test	Tensor	Spatial coordinates of the locations for prediction	required
theta	tuple[float, float, float]	theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.	required
neighbor_size	Optional[int]	The number of nearest neighbors used for NNGP approximation. Default being 20.	20
q	Optional[float]	Confidence coverage for the prediction interval. Default being 0.95.	0.95

Returns:

Name	Type	Description
w_test	Tensor	torch.Tensor The kriging prediction.
pred_U	Tensor	torch.Tensor Confidence upper bound.
pred_L	Tensor	torch.Tensor Confidence lower bound.

See Also

Zhan, Wentao, and Abhirup Datta. 2024. “Neural Networks for Geospatial Data.” Journal of the American Statistical Association, June, 1–21. doi:10.1080/01621459.2024.2356293.

Source code in geospaNN/utils.py

def krig_pred(w_train: torch.Tensor,
              coord_train: torch.Tensor,
              coord_test: torch.Tensor,
              theta: tuple[float, float, float],
              neighbor_size: Optional[int] = 20,
              q: Optional[float] = 0.95
              ) -> torch.Tensor:
    """Kriging prediction (Gaussian process regression) with confidence interval.

    Kriging prediction on testing locations based on the observations on the training locations. The kriging procedure
    assumes the observations are sampled from a Gaussian process, which is paramatrized here to have an exponential covariance
    structure using theta = [sigma^2, phi, tau]. NNGP appriximation is involved for efficient computation of matrix inverse.
    The conditional variance (kriging variance) is used to build the confidence interval using the quantiles (a/2, 1-a/2).
    (see https://arxiv.org/abs/2304.09157, section 4.3 for more details.)

    Parameters:
        w_train:
            Training observations of the spatial random effect without any fixed effect.
        coord_train:
            Spatial coordinates of the training observations.
        coord_test:
            Spatial coordinates of the locations for prediction
        theta:
            theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.
        neighbor_size:
            The number of nearest neighbors used for NNGP approximation. Default being 20.
        q:
            Confidence coverage for the prediction interval. Default being 0.95.

    Returns:
        w_test: torch.Tensor
            The kriging prediction.
        pred_U: torch.Tensor
            Confidence upper bound.
        pred_L: torch.Tensor
            Confidence lower bound.

    See Also:
        Zhan, Wentao, and Abhirup Datta. 2024. “Neural Networks for Geospatial Data.”
        Journal of the American Statistical Association, June, 1–21. doi:10.1080/01621459.2024.2356293.
    """
    sigma_sq, phi, tau = theta
    tau_sq = tau * sigma_sq
    n_test = coord_test.shape[0]

    rank = make_rank(coord_train, neighbor_size, coord_test)

    w_test = torch.zeros(n_test)
    sigma_test = (sigma_sq + tau_sq) * torch.ones(n_test)
    for i in range(n_test):
        ind = rank[i, :]
        cov_sub = make_cov_full(distance(coord_train[ind, :], coord_train[ind, :]), theta, nuggets=True)
        cov_vec = make_cov_full(distance(coord_train[ind, :], coord_test[i, :]), theta, nuggets=False).reshape(-1)
        bi = torch.linalg.solve(cov_sub, cov_vec)
        w_test[i] = torch.dot(bi.T, w_train[ind]).squeeze()
        sigma_test[i] = sigma_test[i] - torch.dot(bi.reshape(-1), cov_vec)
    p = scipy.stats.norm.ppf((1 + q) / 2, loc=0, scale=1)
    sigma_test = torch.sqrt(sigma_test)
    pred_U = w_test + p * sigma_test
    pred_L = w_test - p * sigma_test

    return w_test, pred_U, pred_L

make_bf(coord, theta, neighbor_size=20)

Obtain NNGP approximation with implicit covariance matrix

Find the upper triangular matrix B and diagonal matrix F such that (I-B)'F^{-1}(I-B) appriximate the precision matrix (inverse of the covariance matrix). (see https://arxiv.org/abs/2102.13299 for more details.) Here only coordinates and spatial parameters are needed to represent the exponential covariance implicitly, thus being more memory-efficient than make_bf_from_cov.

Parameters:

Name	Type	Description	Default
coord	Tensor	The nxd covariate array.	required
neighbor_size	Optional[int]	The number of nearest neighbors used used for NNGP approximation. Default being 20.	20
theta	tuple[float, float, float]	theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.	required

Returns:

Name	Type	Description
I_B	Sparse_B	The sparse representation of I-B.
F	Tensor	The vector representation of the diagonal matrix.

See Also

make_bf_from_cov : Obtain NNGP approximation of a covariance matrix Datta, Abhirup. "Sparse nearest neighbor Cholesky matrices in spatial statistics." arXiv preprint arXiv:2102.13299 (2021).

Source code in geospaNN/utils.py

def make_bf(coord: torch.Tensor,  #### could add a make_bf from cov (resolved)
            theta: tuple[float, float, float],
            neighbor_size: Optional[int] = 20,
            ) -> Tuple[Sparse_B, torch.Tensor]:
    """Obtain NNGP approximation with implicit covariance matrix

    Find the upper triangular matrix B and diagonal matrix F such that (I-B)'F^{-1}(I-B) appriximate the precision matrix
    (inverse of the covariance matrix). (see https://arxiv.org/abs/2102.13299 for more details.) Here only coordinates and
    spatial parameters are needed to represent the exponential covariance implicitly,
    thus being more memory-efficient than make_bf_from_cov.

    Parameters:
        coord:
            The nxd covariate array.
        neighbor_size:
            The number of nearest neighbors used used for NNGP approximation. Default being 20.
        theta:
            theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.

    Returns:
        I_B:
            The sparse representation of I-B.
        F:
            The vector representation of the diagonal matrix.

    See Also:
        make_bf_from_cov : Obtain NNGP approximation of a covariance matrix \

        Datta, Abhirup. "Sparse nearest neighbor Cholesky matrices in spatial statistics."
        arXiv preprint arXiv:2102.13299 (2021).
    """
    n = coord.shape[0]
    rank = make_rank(coord, neighbor_size)
    B = torch.zeros((n, neighbor_size))
    ind_list = np.zeros((n, neighbor_size)).astype(int) - 1
    F = torch.zeros(n)
    for i in range(n):
        F[i] = make_cov_full(torch.tensor([0]), theta, nuggets = True)
        ind = rank[i, :][rank[i, :] <= i]
        if len(ind) == 0:
            continue
        cov_sub = make_cov_full(distance(coord[ind, :], coord[ind, :]), theta, nuggets = True)
        if torch.linalg.matrix_rank(cov_sub) == cov_sub.shape[0]:
            cov_vec = make_cov_full(distance(coord[ind, :], coord[i, :]), theta).reshape(-1)
            #### nuggets is not specified since its off-diagonal
            bi = torch.linalg.solve(cov_sub, cov_vec)
            B[i, range(len(ind))] = bi
            ind_list[i, range(len(ind))] = ind
            F[i] = F[i] - torch.inner(cov_vec, bi)

    I_B = Sparse_B(torch.concatenate([torch.ones((n, 1)), -B], axis=1),
                   np.concatenate([np.arange(0, n).reshape(n, 1), ind_list], axis = 1))

    return I_B, F

make_bf_from_cov(cov, neighbor_size)

Obtain NNGP approximation of a covariance matrix

Find the upper triangular matrix B and diagonal matrix F such that (I-B)'F^{-1}(I-B) appriximate the precision matrix (inverse of the covariance matrix). The level of approximation increase with the neighbor size. When using the full neighbor, the NNGP appriximation degrade to the Cholesky decomposition. (see https://arxiv.org/abs/2102.13299 for more details.)

Parameters:

Name	Type	Description	Default
cov	Tensor	The nxn covariance matrix.	required
neighbor_size	int	The number of nearest neighbors used for NNGP approximation.	required

Returns:

Name	Type	Description
I_B	Sparse_B	Sparse_B The sparse representation of I-B.
F	Sparse_B	torch.Tensor The vector representation of the diagonal matrix.

See Also

make_bf : Obtain NNGP approximation with implicit covariance matrix Datta, Abhirup, et al. "Hierarchical nearest-neighbor Gaussian process models for large geostatistical datasets." Journal of the American Statistical Association 111.514 (2016): 800-812.

Datta, Abhirup. "Sparse nearest neighbor Cholesky matrices in spatial statistics." arXiv preprint arXiv:2102.13299 (2021).

Source code in geospaNN/utils.py

def make_bf_from_cov(cov: torch.Tensor,
                     neighbor_size: int,
                     ) -> Sparse_B:
    """Obtain NNGP approximation of a covariance matrix

    Find the upper triangular matrix B and diagonal matrix F such that (I-B)'F^{-1}(I-B) appriximate the precision matrix
    (inverse of the covariance matrix). The level of approximation increase with the neighbor size. When using the full neighbor,
    the NNGP appriximation degrade to the Cholesky decomposition. (see https://arxiv.org/abs/2102.13299 for more details.)

    Parameters:
        cov:
            The nxn covariance matrix.
        neighbor_size:
            The number of nearest neighbors used for NNGP approximation.

    Returns:
        I_B: Sparse_B
            The sparse representation of I-B.
        F: torch.Tensor
            The vector representation of the diagonal matrix.

    See Also:
        make_bf : Obtain NNGP approximation with implicit covariance matrix \

        Datta, Abhirup, et al. "Hierarchical nearest-neighbor Gaussian process models for large geostatistical datasets."
        Journal of the American Statistical Association 111.514 (2016): 800-812.

        Datta, Abhirup. "Sparse nearest neighbor Cholesky matrices in spatial statistics."
        arXiv preprint arXiv:2102.13299 (2021).
    """
    n = cov.shape[0]
    B = torch.zeros((n, neighbor_size))
    ind_list = np.zeros((n, neighbor_size)).astype(int) - 1
    F = torch.zeros(n)
    rank = np.argsort(-cov, axis=-1) #### consider replace using make_rank?
    rank = rank[:, 1:(neighbor_size + 1)]
    for i in range(n):
        F[i] = cov.diag()[i]
        ind = rank[i, :][rank[i, :] <= i]
        if len(ind) == 0:
            continue
        cov_sub = cov[ind.reshape(-1, 1), ind.reshape(1, -1)]
        if torch.linalg.matrix_rank(cov_sub) == cov_sub.shape[0]:
            cov_vec = cov[ind, i].reshape(-1)
            bi = torch.linalg.solve(cov_sub, cov_vec)
            B[i, range(len(ind))] = bi
            ind_list[i, range(len(ind))] = ind
            F[i] = F[i] - torch.inner(cov_vec, bi)

    I_B = Sparse_B(torch.concatenate([torch.ones((n, 1)), -B], axis=1),
                   np.concatenate([np.arange(0, n).reshape(n, 1), ind_list], axis=1))

    return I_B, F #### Shall we return NNGP_cov instead?

make_cov(coord, theta, NNGP=True, neighbor_size=20)

Compose covariance matrix.

Compose a covariance matrix in the exponential covariance family using the coordinates and spatial parameters. NNGP approximation is introduced for efficient representation. (see https://arxiv.org/abs/2102.13299 for more details.)

Parameters:

Name	Type	Description	Default
coord	Tensor	The nxd covariate array.	required
theta	tuple[float, float, float]	theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.	required
NNGP	Optional[bool]	Whether use NNGP approximation (recommended and used by default).	True
neighbor_size	int	Number of nearest neighbors used for NNGP approximation, default value is 20.	20

Returns:

Name	Type	Description
cov	Tensor	A covariance matrix as torch.Tensor (dense representation) or NNGP_cov (sparse representation).

See Also

Datta, Abhirup. "Sparse nearest neighbor Cholesky matrices in spatial statistics." arXiv preprint arXiv:2102.13299 (2021).

Source code in geospaNN/utils.py

def make_cov(coord: torch.Tensor,
             theta: tuple[float, float, float],
             NNGP: Optional[bool] = True,
             neighbor_size: int = 20
             ) -> torch.Tensor:
    """Compose covariance matrix.

    Compose a covariance matrix in the exponential covariance family using the coordinates and spatial parameters.
    NNGP approximation is introduced for efficient representation. (see https://arxiv.org/abs/2102.13299 for more details.)

    Parameters:
        coord:
            The nxd covariate array.
        theta:
            theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.
        NNGP:
            Whether use NNGP approximation (recommended and used by default).
        neighbor_size:
            Number of nearest neighbors used for NNGP approximation, default value is 20.

    Returns:
        cov:
            A covariance matrix as torch.Tensor (dense representation) or NNGP_cov (sparse representation).

    See Also:
        Datta, Abhirup. "Sparse nearest neighbor Cholesky matrices in spatial statistics."
        arXiv preprint arXiv:2102.13299 (2021).
    """
    if not NNGP:
        dist = distance(coord, coord)
        cov = make_cov_full(dist, theta, nuggets = True) #### could add a make_bf from cov (resolved)
        return cov
    else:
        I_B, F_diag = make_bf(coord, theta, neighbor_size) #### could merge into one step
        cov = NNGP_cov(I_B.B, F_diag, I_B.Ind_list)
        return cov

make_cov_full(dist, theta, nuggets=False)

Compose covariance matrix from the distance matrix with dense representation.

Compose a covariance matrix in the exponential covariance family (other options to be implemented) from the distance matrix. The returned object class depends on the input distance matrix.

Parameters:

Name	Type	Description	Default
dist	Tensor | ndarray	The nxn distance matrix	required
theta	tuple[float, float, float]	theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.	required
nuggets	Optional[bool]	Whether to include nuggets term in the covariance matrix (added to the diagonal).	False

Returns:

Name	Type	Description
cov	Tensor | ndarray	A covariance matrix.

Source code in geospaNN/utils.py

def make_cov_full(dist: torch.Tensor | np.ndarray,
                  theta: tuple[float, float, float],
                  nuggets: Optional[bool] = False,
                  ) -> torch.Tensor | np.ndarray:
    """Compose covariance matrix from the distance matrix with dense representation.

    Compose a covariance matrix in the exponential covariance family (other options to be implemented) from the distance
    matrix. The returned object class depends on the input distance matrix.

    Parameters:
        dist:
            The nxn distance matrix
        theta:
            theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.
        nuggets:
            Whether to include nuggets term in the covariance matrix (added to the diagonal).

    Returns:
        cov:
            A covariance matrix.
    """
    sigma_sq, phi, tau = theta
    tau_sq = tau * sigma_sq
    if isinstance(dist, float) or isinstance(dist, int):
        dist = torch.Tensor(dist)
        n = 1
    else:
        n = dist.shape[-1]
    if isinstance(dist, torch.Tensor):
        cov = sigma_sq * torch.exp(-phi * dist)
    else:
        cov = sigma_sq * np.exp(-phi * dist)
    if nuggets:
        shape_temp = list(cov.shape)[:-2] + [1 ,1]
        if isinstance(dist, torch.Tensor):
            cov += tau_sq * torch.eye(n).repeat(*shape_temp).squeeze()
        else:
            cov += tau_sq * np.eye(n).squeeze() #### need improvement
    return cov

make_graph(X, Y, coord, neighbor_size=20, Ind_list=None)

Compose the data with graph information to work on.

This function connects each node to its nearest neighbors and records the edge indexes in two forms for the downstream graph operations.

Parameters:

Name	Type	Description	Default
X	Tensor	nxp array of the covariates.	required
Y	Tensor	Length n vector as the response.	required
coord	Tensor	nxd array of the coordinates	required
neighbor_size	Optional[int]	The number of nearest neighbors used for NNGP approximation. Default being 20.	20
Ind_list	Optional	An optional index list. If provided, ommit the make_rank() step in the function.	None

Returns:

Name	Type	Description
data	Data	torch_geometric.data.Data Data that can be processed by the torch_geometric framework. data.x contains the covariates array, data.y contains the response vector, data.pos contains the spatial coordinates, data.edge_index contains the indexes of form [i,j] where location j is in the nearest neighbor of location i. data.edge_attr contains the concatenated numbering of the neighbors. For each location, the numbering is of the form [0, 1, ... , k] where k is the number of the nearest neighbors. This attribute is mainly used in messaging passing.

See Also: make_rank : Compose the nearest neighbor index list based on the coordinates. torch_geometric.data.Data : A data object describing a homogeneous graph.

Source code in geospaNN/utils.py

def make_graph(X: torch.Tensor,
               Y: torch.Tensor,
               coord: torch.Tensor,
               neighbor_size: Optional[int] = 20,
               Ind_list: Optional = None
               ) -> torch_geometric.data.Data:
    """Compose the data with graph information to work on.

    This function connects each node to its nearest neighbors and records the edge indexes in two forms for the downstream
    graph operations.

    Parameters:
        X:
            nxp array of the covariates.
        Y:
            Length n vector as the response.
        coord:
            nxd array of the coordinates
        neighbor_size:
            The number of nearest neighbors used for NNGP approximation. Default being 20.
        Ind_list:
            An optional index list. If provided, ommit the make_rank() step in the function.

    Returns:
        data: torch_geometric.data.Data \
            Data that can be processed by the torch_geometric framework.\
            data.x contains the covariates array,\
            data.y contains the response vector,\
            data.pos contains the spatial coordinates,\
            data.edge_index contains the indexes of form [i,j] where location j is in the nearest neighbor of location i.\
            data.edge_attr contains the concatenated numbering of the neighbors.\
            For each location, the numbering is of the form [0, 1, ... , k] where k is the number of the nearest neighbors.\
            This attribute is mainly used in messaging passing.

    See Also:
    make_rank : Compose the nearest neighbor index list based on the coordinates. \
    torch_geometric.data.Data : A data object describing a homogeneous graph.
    """
    n = X.shape[0]
    # Compute the edges of the graph
    edges = []
    neighbor_idc = []
    # Initialize the edges, the edges are predefined for the first m + 1 points
    # Find the m nearest neighbors for each remaining point
    if Ind_list is None:
        Ind_list = make_rank(coord, neighbor_size)
    for i in range(1, n):
        for j, idx in enumerate(Ind_list[i]):
            if idx < i:
                edges.append([idx, i])
                neighbor_idc.append(j)
            elif j >= i:
                break

    edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()
    edge_attr = torch.tensor(neighbor_idc).reshape(-1, 1)  # denotes the index of the neighbor
    data = torch_geometric.data.Data(x=X, y=Y, pos=coord, edge_index=edge_index, edge_attr=edge_attr)
    assert data.validate(raise_on_error=True)
    return data

make_rank(coord, neighbor_size, coord_ref=None)

Compose the nearest neighbor index list based on the coordinates.

Find the indexes of nearest neighbors in reference set for each location i in the main set. The index is based on the increasing order of the distances between ith location and the locations in the reference set.

Parameters:

Name	Type	Description	Default
coord	Tensor	The nxd coordinates array of target locations.	required
neighbor_size	int		required
coord_ref	Optional	The n_refxd coordinates array of reference locations. If None, use the target set itself as the reference. (Any location's neighbor does not include itself.)	None

Returns:

Name	Type	Description
rank_list	ndarray	A nxp array. The ith row is the indexes of the nearest neighbors for the ith location, ordered by the distance.

Source code in geospaNN/utils.py

def make_rank(coord: torch.Tensor,
              neighbor_size: int,
              coord_ref: Optional = None
              ) -> np.ndarray:
    """Compose the nearest neighbor index list based on the coordinates.

    Find the indexes of nearest neighbors in reference set for each location i in the main set.
    The index is based on the increasing order of the distances between ith location and the locations in the reference set.

    Parameters:
        coord:
            The nxd coordinates array of target locations.
        neighbor_size:
        `   Suppose neighbor_size = k, only the top k-nearest indexes will be returned.
        coord_ref:
            The n_refxd coordinates array of reference locations. If None, use the target set itself as the reference.
            (Any location's neighbor does not include itself.)

    Returns:
        rank_list:
            A nxp array. The ith row is the indexes of the nearest neighbors for the ith location, ordered by the distance.
    """
    if coord_ref is None: neighbor_size += 1
    knn = NearestNeighbors(n_neighbors=neighbor_size)
    knn.fit(coord)
    if coord_ref is None:
        coord_ref = coord
        rank = knn.kneighbors(coord_ref)[1]
        return rank[:, 1:]
    else:
        rank = knn.kneighbors(coord_ref)[1]
        return rank[:, 0:]

rmvn(mu, cov, sparse=True)

Randomly generate sample from multivariate normal distribution

Generate random sample from a multivariate normal distribution with specified mean and covariance.

Parameters:

Name	Type	Description	Default
mu	Tensor	The additional mean of multivariate normal of length n.	required
cov	Tensor | NNGP_cov		required
sparse	Optional[bool]	Designed for sparse representation, not implemented yet.	True

Returns:

Name	Type	Description
sample	Tensor	A random sample from the multivariate normal distribution.

Source code in geospaNN/utils.py

def rmvn(mu: torch.Tensor,
         cov: torch.Tensor | NNGP_cov,
         sparse: Optional[bool] = True) \
        -> torch.Tensor:
    """Randomly generate sample from multivariate normal distribution

    Generate random sample from a multivariate normal distribution with specified mean and covariance.

    Parameters:
        mu:
            The additional mean of multivariate normal of length n.
        cov:
        `   The nxn covariance matrix. When use torch.Tensor for the dense representation, Cholesky's decomposition is used
            for correlating the i.i.d normal sample. When use NNGP_cov object for sparse representation, use NNGP to approximate
            the correlating process. Dense representation is not recommended for large sample size.
        sparse:
            Designed for sparse representation, not implemented yet.

    Returns:
        sample:
            A random sample from the multivariate normal distribution.
    """
    n = len(mu) #### Check dimensionality
    if isinstance(cov, torch.Tensor):
        if n >= 2000: warnings.warn("Too large for cholesky decomposition, please try to use NNGP")
        D = torch.linalg.cholesky(cov)
        res = torch.matmul(torch.randn(1, n), D.t()) + mu
    elif isinstance(cov, NNGP_cov):
        if sparse:
            res = cov.correlate(torch.randn(1, n).reshape(-1)) + mu
        else:
            warnings.warn("To be implemented.")
    else:
        warnings.warn("Covariance matrix should be in the format of torch.Tensor or NNGP_cov.")
        return

    return  res.reshape(-1)

spatial_order(X, Y, coord, method='max-min', numpropose=2, seed=2024)

Spatial ordering for data

Order the data according to their spatial locations. A spatial ordering is necessary for NNGP to represent a valid spatial process. (Datta et.al 2016) Method 'coord-sum' stands for a simple spatial ordering by the summation of coordinates. Method 'max-min' stands for the max-min ordering based on Euclidean distance. "Basically, this ordering starts at a point in the center, then adds points one at a time that maximize the minimum distance from all previous points in the ordering." (Katzfuss & Guinness 2021)

Parameters:

Name	Type	Description	Default
X	Tensor	nxp array of the covariates.	required
Y	Tensor	Length n vector as the response.	required
coord	Tensor	nxd array of the coordinates	required
method	str	Method 'coord-sum' stands for a simple spatial ordering by the summation of coordinates. Method 'max-min' stands for the max-min ordering based on Euclidean distance. (Katzfuss & Guinness 2021)	'max-min'

Returns:

Type	Description
Tuple[Tensor, Tensor, Tensor, Tensor]	Ordered X, Ordered Y, Ordered coordinates, order

See Also

Datta, Abhirup, et al. "Hierarchical nearest-neighbor Gaussian process models for large geostatistical datasets." Journal of the American Statistical Association 111.514 (2016): 800-812. Katzfuss, Matthias & Guinness, Joseph. "A General Framework for Vecchia Approximations of Gaussian Processes." Statist. Sci. 36 (1) 124 - 141, February 2021. https://doi.org/10.1214/19-STS755

Source code in geospaNN/utils.py

def spatial_order(X: torch.Tensor,
                  Y: torch.Tensor,
                  coord: torch.Tensor,
                  method: str = 'max-min',
                  numpropose: Optional[int] = 2,
                  seed: Optional[int] = 2024
                  ) -> Tuple[torch.Tensor,torch.Tensor,torch.Tensor,torch.Tensor]:
    """Spatial ordering for data

    Order the data according to their spatial locations. A spatial ordering is necessary for NNGP to represent a valid
    spatial process. (Datta et.al 2016)
    Method 'coord-sum' stands for a simple spatial ordering by the summation of coordinates.
    Method 'max-min' stands for the max-min ordering based on Euclidean distance. "Basically, this ordering starts at a
    point in the center, then adds points one at a time that maximize the minimum distance from all previous points in
    the ordering." (Katzfuss & Guinness 2021)

    Parameters:
        X:
            nxp array of the covariates.
        Y:
            Length n vector as the response.
        coord:
            nxd array of the coordinates
        method:
            Method 'coord-sum' stands for a simple spatial ordering by the summation of coordinates.
            Method 'max-min' stands for the max-min ordering based on Euclidean distance. (Katzfuss & Guinness 2021)

    Returns:
        Ordered X, Ordered Y, Ordered coordinates, order

    See Also:
        Datta, Abhirup, et al. "Hierarchical nearest-neighbor Gaussian process models for large geostatistical datasets."
        Journal of the American Statistical Association 111.514 (2016): 800-812. \
        Katzfuss, Matthias & Guinness, Joseph. "A General Framework for Vecchia Approximations of Gaussian Processes."
        Statist. Sci. 36 (1) 124 - 141, February 2021. https://doi.org/10.1214/19-STS755
    """
    n = coord.shape[0]
    if n >= 10000 and method == 'max-min':
        warnings.warn("Too large for max-min ordering, switch to 'coord-sum' ordering!")
        method = 'coord-sum'

    d = coord.shape[1]
    random.seed(seed)
    if method == 'max-min':
        remaininginds = list(range(n))
        orderinds = torch.zeros(n, dtype=torch.int)
        distmp = distance(coord, coord.mean(dim = 0)).reshape(-1)
        ordermp = torch.argsort(distmp)
        orderinds[0] = ordermp[0]
        remaininginds.remove(orderinds[0])
        for j in range(1,n):
            randinds = random.sample(remaininginds, min(numpropose, len(remaininginds)))
            distarray = distance(coord[orderinds[0:j],:], coord[torch.tensor(randinds),:])
            bestind = torch.argmax(distarray.min(axis = 1).values)
            orderinds[j] = randinds[bestind]
            remaininginds.remove(orderinds[j])
    elif method == 'coord-sum':
        orderinds = torch.argsort(coord.sum(axis = 1))
    else:
        warnings.warn("Keep the order")
        orderinds = torch.tensor(range(n))

    return X[orderinds,:], Y[orderinds], coord[orderinds,:], orderinds

split_data(X, Y, coord, neighbor_size=20, val_proportion=0.2, test_proportion=0.2)

Split the data into training, validation, and testing parts and add the graph information respectively.

This function split the data with a user specified proportions and build the graph information. The output are the data sets that can be directly processed within the torch_geomtric framework.

Parameters:

Name	Type	Description	Default
X	Tensor	nxp array of the covariates.	required
Y	Tensor	Length n vector as the response.	required
coord	Tensor	nxd array of the coordinates	required
neighbor_size	Optional[int]	The number of nearest neighbors used for NNGP approximation. Default being 20.	20
val_proportion	float	The proportion of training set splitted as validation set.	0.2
test_proportion	float	The proportion of whole data splitted as testing set.	0.2

Returns:

Type	Description
tuple[Data, Data, Data]	data_train, data_val, data_test

See Also

make_graph : Compose the data with graph information to work on. torch_geometric.data.Data : A data object describing a homogeneous graph.

Source code in geospaNN/utils.py

def split_data(X: torch.Tensor,
               Y: torch.Tensor,
               coord: torch.Tensor,
               neighbor_size: Optional[int] = 20,
               val_proportion: float = 0.2,
               test_proportion: float = 0.2
               ) -> tuple[torch_geometric.data.Data, torch_geometric.data.Data, torch_geometric.data.Data]:
    """Split the data into training, validation, and testing parts and add the graph information respectively.

    This function split the data with a user specified proportions and build the graph information. The output are the
    data sets that can be directly processed within the torch_geomtric framework.

    Parameters:
        X:
            nxp array of the covariates.
        Y:
            Length n vector as the response.
        coord:
            nxd array of the coordinates
        neighbor_size:
            The number of nearest neighbors used for NNGP approximation. Default being 20.
        val_proportion:
            The proportion of training set splitted as validation set.
        test_proportion:
            The proportion of whole data splitted as testing set.

    Returns:
        data_train, data_val, data_test

    See Also:
        make_graph : Compose the data with graph information to work on. \
        torch_geometric.data.Data : A data object describing a homogeneous graph.
    """
    X_train_val, X_test, Y_train_val, Y_test, coord_train_val, coord_test = train_test_split(
        X, Y, coord, test_size = test_proportion
    )

    X_train, X_val, Y_train, Y_val, coord_train, coord_val = train_test_split(
        X_train_val, Y_train_val, coord_train_val, test_size = val_proportion/(1 - test_proportion)
    )

    data_train = make_graph(X_train, Y_train, coord_train, neighbor_size)
    data_val = make_graph(X_val, Y_val, coord_val, neighbor_size)
    data_test = make_graph(X_test, Y_test, coord_test, neighbor_size)

    return data_train, data_val, data_test

split_loader(data, batch_size=None)

Create mini-batches for GNN training

This functions further split a data for mini-batch training of GNNs on large-scale graphs where full-batch training is not feasible. Note that only source nodes are splited into batches, each batch will contain edges originate from those source nodes. There might be interactions among the target nodes across batches.

Parameters:

Name	Type	Description	Default
data	Data	Data with graph information (output of split_data() or make_graph()).	required
batch_size	Optional[int]	Size of mini-batches, default value being n/20.	None

Returns:

Name	Type	Description
loader	DataLoaders	A dataloader that can be enumerated for mini-batch training.

See Also

make_graph : Compose the data with graph information to work on. split_data : Split the data into training, validation, and testing parts and add the graph information respectively. torch_geometric.loader : A data loader that performs neighbor sampling as introduced in the

Hamilton, Will, Zhitao Ying, and Jure Leskovec. "Inductive representation learning on large graphs." Advances in neural information processing systems 30 (2017).

Source code in geospaNN/utils.py

def split_loader(data: torch_geometric.data.Data,
                 batch_size: Optional[int] = None
                 ) -> torch.DataLoaders:
    """Create mini-batches for GNN training

    This functions further split a data for mini-batch training of GNNs on large-scale graphs
    where full-batch training is not feasible. Note that only source nodes are splited into batches,
    each batch will contain edges originate from those source nodes.
    There might be interactions among the target nodes across batches.

    Parameters:
        data:
            Data with graph information (output of split_data() or make_graph()).
        batch_size:
            Size of mini-batches, default value being n/20.

    Returns:
        loader:
            A dataloader that can be enumerated for mini-batch training.

    See Also:
        make_graph : Compose the data with graph information to work on. \
        split_data : Split the data into training, validation, and testing parts and add the graph information respectively. \
        torch_geometric.loader : A data loader that performs neighbor sampling as introduced in the \

    Hamilton, Will, Zhitao Ying, and Jure Leskovec. "Inductive representation learning on large graphs."
    Advances in neural information processing systems 30 (2017).
    """
    if batch_size is None:
        batch_size  = int(data.x.shape[0]/20)
    loader = NeighborLoader(data,
                            input_nodes=torch.tensor(range(data.x.shape[0])),num_neighbors=[-1],
                            batch_size=batch_size, replace=False, shuffle=True)
    return loader

theta_update(w, coord, theta0=None, neighbor_size=20, BRISC=True, min_tau=0.001)

Update the spatial parameters using maximum likelihood.

This function updates the spatial parameters by assuming the observations are from a Gaussian Process with exponential covariance function. Spatial coordinates and initial values of theta are input for building the covariance. By default, L-BFGS-B algorithm is used to optimize the likelihood. Since the likelihood is computed repeatedly, NNGP approximation is used for efficient computation of the log-likelihood, with a default neighbor size being 20. Note that we currently use the Cpp implementation of L-BFGS-B from the R-package BRISC (Saha & Datta 2018) as default, which is faster than the python version. In the future, we will introduce a python wrapper for the Cpp implementation to get free of R component.

Parameters:

Name	Type	Description	Default
theta0	Optional[Tensor]	Initial values of the spatial parameters. theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.	None
w	Tensor	Length n observations of the spatial random effect without any fixed effect.	required
coord	Tensor	nx2 spatial coordinates of the observations.	required
neighbor_size	Optional[int]	The number of nearest neighbors used for NNGP approximation. Default being 20.	20
BRISC	Optional[bool]	Whether to use the optimization from BRISC. Default being True. Setting as False will largely increase the running time.	True
min_tau	Optional[float]	A minimum value of tau to avoid sinularity issue in matrix inversion. Default being 0.001.	0.001

Returns:

Name	Type	Description
theta_updated	array	An updated tuple of the spatial paramaters.

See Also

Datta, Abhirup, et al. "Hierarchical nearest-neighbor Gaussian process models for large geostatistical datasets." Journal of the American Statistical Association 111.514 (2016): 800-812.

Zhu, Ciyou, et al. "Algorithm 778: L-BFGS-B: Fortran subroutines for large-scale bound-constrained optimization." ACM Transactions on mathematical software (TOMS) 23.4 (1997): 550-560.

Saha, Arkajyoti, and Abhirup Datta. "BRISC: bootstrap for rapid inference on spatial covariances." Stat 7.1 (2018): e184.

Source code in geospaNN/utils.py

def theta_update(w: torch.Tensor,
                 coord: torch.Tensor,
                 theta0: Optional[torch.Tensor] = None,
                 neighbor_size: Optional[int] = 20,
                 BRISC: Optional[bool] = True,
                 min_tau: Optional[float] = 0.001
                 ) -> np.array:
    """Update the spatial parameters using maximum likelihood.

    This function updates the spatial parameters by assuming the observations are from a Gaussian Process with exponential
    covariance function. Spatial coordinates and initial values of theta are input for building the covariance.
    By default, L-BFGS-B algorithm is used to optimize the likelihood.
    Since the likelihood is computed repeatedly, NNGP approximation is used for efficient computation of the log-likelihood,
    with a default neighbor size being 20.
    Note that we currently use the Cpp implementation of L-BFGS-B from the R-package BRISC (Saha & Datta 2018) as default,
    which is faster than the python version.
    In the future, we will introduce a python wrapper for the Cpp implementation to get free of R component.

    Parameters:
        theta0:
            Initial values of the spatial parameters.
            theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.
        w:
            Length n observations of the spatial random effect without any fixed effect.
        coord:
            nx2 spatial coordinates of the observations.
        neighbor_size:
            The number of nearest neighbors used for NNGP approximation. Default being 20.
        BRISC:
            Whether to use the optimization from BRISC. Default being True. Setting as False will largely increase the running time.
        min_tau:
            A minimum value of tau to avoid sinularity issue in matrix inversion. Default being 0.001.

    Returns:
        theta_updated:
            An updated tuple of the spatial paramaters.

    See Also:
        Datta, Abhirup, et al. "Hierarchical nearest-neighbor Gaussian process models for large geostatistical datasets."
        Journal of the American Statistical Association 111.514 (2016): 800-812.

        Zhu, Ciyou, et al. "Algorithm 778: L-BFGS-B: Fortran subroutines for large-scale bound-constrained optimization."
        ACM Transactions on mathematical software (TOMS) 23.4 (1997): 550-560.

        Saha, Arkajyoti, and Abhirup Datta. "BRISC: bootstrap for rapid inference on spatial covariances."
        Stat 7.1 (2018): e184.
    """
    warnings.filterwarnings("ignore")
    w = w.detach().numpy()
    coord = coord.detach().numpy()

    if not BRISC:
        theta = theta0.detach().numpy()
        if theta0 is None:
            warnings.warn("Theta not initialized, start from [1,1,1].")
            theta = np.array([1,1,1])
        else:
            print('Theta updated from')
            print(theta)
        n_train = w.shape[0]
        rank = make_rank(coord, neighbor_size)
        if n_train <= 2000:
            dist = distance_np(coord, coord)

            def likelihood(theta):
                sigma, phi, tau = theta
                cov = sigma * (np.exp(-phi * dist)) + tau * sigma * np.eye(n_train)  # need dist, n

                term1 = 0
                term2 = 0
                for i in range(n_train):
                    ind = rank[i, :][rank[i, :] <= i]
                    id = np.append(ind, i)

                    sub_cov = cov[ind, :][:, ind]
                    if np.linalg.det(sub_cov):
                        bi = np.linalg.solve(cov[ind, :][:, ind], cov[ind, i])
                    else:
                        bi = np.zeros(ind.shape)
                    I_B_i = np.append(-bi, 1)
                    F_i = cov[i, i] - np.inner(cov[ind, i], bi)
                    decor_res = np.sqrt(np.reciprocal(F_i)) * np.dot(I_B_i, w[id])
                    term1 += np.log(F_i)
                    term2 += decor_res ** 2
                return (term1 + term2)

        else:
            def likelihood(theta):
                sigma, phi, tau = theta

                term1 = 0
                term2 = 0
                for i in range(n_train):
                    ind = rank[i, :][rank[i, :] <= i]
                    if len(ind) == 0:
                        F_i = (1 + tau) * sigma
                        term1 += np.log(F_i)
                        term2 += w[i] ** 2 / F_i
                        continue

                    id = np.append(ind, i)

                    sub_dist = distance_np(coord[ind, :], coord[ind, :])
                    sub_dist_vec = distance_np(coord[ind, :], coord[i, :].reshape(1, -1)).reshape(-1)
                    sub_cov = sigma * (np.exp(-phi * sub_dist)) + tau * sigma * np.eye(len(ind))
                    sub_cov_vec = sigma * (np.exp(-phi * sub_dist_vec))
                    if np.linalg.det(sub_cov):
                        bi = np.linalg.solve(sub_cov, sub_cov_vec)
                    else:
                        bi = np.zeros(ind.shape)
                    I_B_i = np.append(-bi, 1)
                    F_i = (1 + tau) * sigma - np.inner(sub_cov_vec, bi)
                    decor_res = np.sqrt(np.reciprocal(F_i)) * np.dot(I_B_i, w[id])
                    term1 += np.log(F_i)
                    term2 += decor_res ** 2
                return (term1 + term2)

        res = scipy.optimize.minimize(likelihood, theta, method = 'L-BFGS-B',
                                      bounds = [(0, None), (0, None), (min_tau, None)])
        print('Theta estimated as')
        print(res.x)
        return res.x

    else:
        _, theta = BRISC_estimation(w, None, coord)
        theta[2] = max(theta[2], min_tau)
        print('Theta estimated as')
        print(theta)
        return theta

plot_PDP(model, X, names=[], save_path='./', save=False, split=False)

Partial dependency plot for model on the data.

A Partial Dependence Plot (PDP) is a visualization tool used to illustrate the relationship between a selected feature and the predicted outcome of a machine learning model, while averaging out the effects of other features. This helps to understand the marginal influence of a single feature on the model's predictions in a more interpretable way.

Parameters:

Name	Type	Description	Default
model		Usually a model in nngls class. Can take model with .() method that take tensor X as input and predicted scalar value Y as output. (to implement for more models)	required
X	tensor	nxp array of the covariates.	required
names	Optional[list]	List of names for variable, if not specified, use "variable 1" to "variable p".	[]
save_path	Optional[str]	Directory to save the plot. Defaults to "./".	'./'
save	bool	Whether to save the PDPs to the working directory. Default False.	False
split	bool	Whether to return the PDPs as a list of PartialDependenceDisplay object for single variabls or a whole	False

Returns:

Type	Description
	A sklearn.inspection.PartialDependenceDisplay object contains PDPs for each variable.

Source code in geospaNN/visualize.py

def plot_PDP(model,
             X: torch.tensor,
             names: Optional[list] = [],
             save_path: Optional[str] = "./",
             save: bool = False,
             split: bool = False):
    """Partial dependency plot for model on the data.

    A Partial Dependence Plot (PDP) is a visualization tool used to illustrate the relationship between a selected feature
    and the predicted outcome of a machine learning model, while averaging out the effects of other features.
    This helps to understand the marginal influence of a single feature on the model's predictions in a more interpretable way.

    Parameters:
        model:
            Usually a model in nngls class. Can take model with .() method that take tensor X as input and
            predicted scalar value Y as output. (to implement for more models)
        X:
            nxp array of the covariates.
        names:
            List of names for variable, if not specified, use "variable 1" to "variable p".
        save_path:
            Directory to save the plot. Defaults to "./".
        save:
            Whether to save the PDPs to the working directory. Default False.
        split:
            Whether to return the PDPs as a list of PartialDependenceDisplay object for single variabls or a whole

    Returns:
        A sklearn.inspection.PartialDependenceDisplay object contains PDPs for each variable.
    """
    X = X.detach().numpy()
    p = X.shape[1]
    Est = _PDP_estimator()
    Est.fit(X, model)
    if len(names) != p:
        warnings.warn("length of names does not match columns of X, replace by variable index")
        names = [f"variable {i + 1}" for i in range(p + 1)]

    if not split:
        figures = PartialDependenceDisplay.from_estimator(estimator=Est, X=X, features=[i for i in range(p)],
                                                          feature_names=names,
                                                          percentiles=(0.05, 0.95))
        if save:
            plt.savefig(save_path + names[0] + ".png")

    else:
        figures = []
        for k in range(p):
            res = PartialDependenceDisplay.from_estimator(estimator=Est, X=X, features=[k],
                                                          feature_names=names[k],
                                                          percentiles=(0.05, 0.95))
            plt.subplots_adjust(left=0.1, bottom=0.1, right=0.9, top=0.9,
                                wspace=0.4, hspace=0.4)
            if save:
                plt.savefig(save_path + names[k] + ".png")
            figures.append(res)

    return figures

plot_PDP_list(model_list, model_names, X, feature_names=[], save_path='./', save=False, split=True)

Partial dependency plot for model on the data.

A Partial Dependence Plot (PDP) is a visualization tool used to illustrate the relationship between a selected feature and the predicted outcome of a machine learning model, while averaging out the effects of other features. This helps to understand the marginal influence of a single feature on the model's predictions in a more interpretable way.

Parameters:

Name	Type	Description	Default
model_list	list	A list of models described in plot_PDP.	required
model_names	list	A list of model names for PDP visualization, expect the same length as model_list.	required
X	tensor	nxp array of the covariates for PDP integration.	required
feature_names	Optional[list]	List of names for variable, if not specified, use "variable 1" to "variable p".	[]
save_path	Optional[str]	Directory to save the plot. Defaults to "./".	'./'
save	bool	Whether to save the PDPs to the working directory. Default False.	False
split	bool	Whether to save the PDP's for different features seperately or as one figure. The default is True.	True

Returns:

Type	Description
	A sklearn.inspection.PartialDependenceDisplay object contains PDPs for each variable.

Source code in geospaNN/visualize.py

def plot_PDP_list(model_list: list,
                  model_names: list,
                  X: torch.tensor,
                  feature_names: Optional[list] = [],
                  save_path: Optional[str] = "./",
                  save: bool = False,
                  split: bool = True):
    """Partial dependency plot for model on the data.

    A Partial Dependence Plot (PDP) is a visualization tool used to illustrate the relationship between a selected feature
    and the predicted outcome of a machine learning model, while averaging out the effects of other features.
    This helps to understand the marginal influence of a single feature on the model's predictions in a more interpretable way.

    Parameters:
        model_list:
            A list of models described in plot_PDP.
        model_names:
            A list of model names for PDP visualization, expect the same length as model_list.
        X:
            nxp array of the covariates for PDP integration.
        feature_names:
            List of names for variable, if not specified, use "variable 1" to "variable p".
        save_path:
            Directory to save the plot. Defaults to "./".
        save:
            Whether to save the PDPs to the working directory. Default False.
        split:
            Whether to save the PDP's for different features seperately or as one figure. The default is True.

    Returns:
        A sklearn.inspection.PartialDependenceDisplay object contains PDPs for each variable.
    """
    l = min(len(model_list), len(model_names))
    if len(model_list) != len(model_names):
        warnings.warn("Number of models different from number of model_names! Use top " + l + " ones.")
    PDP_list = [plot_PDP(model_list[i], X, feature_names) for i in range(l)]
    p = X.shape[1]
    colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
    fig, axes = plt.subplots(p, 1, figsize=(6, 4 * p))  # Adjust height proportionally

    for i in range(l):
        PDP_list[i].plot(ax=axes, line_kw={"label": model_names[i], "color": colors[i]})

    if split:
        for i, ax in enumerate(axes):
            ax.set_title(f"Partial Dependence for Feature {i + 1}")
            ax.legend()  # Ensure legends show up for each axis

            # Save each axis as a separate figure
            fig_single, ax_single = plt.subplots(figsize=(6, 4))  # Create a new figure
            for line in ax.get_lines():  # Copy lines from the original axis
                ax_single.plot(line.get_xdata(), line.get_ydata(), label=line.get_label())

            ax_single.set_title(ax.get_title())
            ax_single.set_xlabel(ax.get_xlabel())
            ax_single.set_ylabel(ax.get_ylabel())
            ax_single.legend()

            # Save the figure for the current axis
            if save:
                fig_single.savefig(save_path + f"PDP_feature_{i + 1}.png")
            plt.close(fig_single)
    else:
        axes.set_title(f"Partial Dependence for Features")
        axes.legend()  # Ensure legends show up for each axis
        fig_single.savefig(save_path + f"PDP_features.png")

plot_log(training_log, theta, save_path='./', save=False)

Output visualization for NN-GLS training.

This is a simple visualization of the training log from geospaNN.nngls_train.train()

Parameters:

Name	Type	Description	Default
training_log	list	A list contains vectors of validation loss, spatial parameters sigma, phi, and tau. Lengths of the vectors must be the same.	required
theta	tuple[float, float, float]	theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.	required
save_path	Optional[str]	Directory to save the plot. Defaults to "./".	'./'
save	bool	Whether to save the PDPs to the working directory. Default False.	False

Source code in geospaNN/visualize.py

def plot_log(training_log: list,
             theta: tuple[float, float, float],
             save_path: Optional[str] = "./",
             save: bool = False):
    """Output visualization for NN-GLS training.

    This is a simple visualization of the training log from geospaNN.nngls_train.train()

    Parameters:
        training_log:
                A list contains vectors of validation loss, spatial parameters sigma, phi, and tau.
                Lengths of the vectors must be the same.
        theta:
            theta[0], theta[1], theta[2] represent sigma^2, phi, tau in the exponential covariance family.
        save_path:
            Directory to save the plot. Defaults to "./".
        save:
            Whether to save the PDPs to the working directory. Default False.
    """
    epoch = len(training_log["val_loss"])
    training_log["epoch"] = list(range(1, epoch + 1))
    training_log = pd.DataFrame(training_log)

    # Melting the dataframe to make it suitable for seaborn plotting
    training_log_melted = training_log[["epoch", "val_loss"]].melt(id_vars='epoch', var_name='Variable', value_name='Value')
    # Plotting with seaborn
    # Creating two subplots side by side
    fig, axes = plt.subplots(1, 2, figsize=(20, 6))
    sns.lineplot(ax=axes[0], data=training_log_melted, x='epoch', y='Value', hue='Variable', style='Variable', markers=False, dashes=False)

    axes[0].set_title('Validation and prediction loss over Epochs (Log Scale) with Benchmark', fontsize=14)
    axes[0].set_xlabel('Epoch', fontsize=15)
    axes[0].set_ylabel('Value (Log Scale)', fontsize=15)
    axes[0].set_yscale('log')
    axes[0].legend(prop={'size': 15})
    axes[0].tick_params(labelsize=14)
    axes[0].grid(True)

    # Second plot (sigma, phi, tau)
    kernel_params_melted = training_log[["epoch", "sigma", "phi", "tau"]].melt(id_vars='epoch', var_name='Variable', value_name='Value')
    ground_truth = {'sigma': theta[0], 'phi': theta[1], 'tau': theta[2]}
    sns.lineplot(ax=axes[1], data=kernel_params_melted, x='epoch', y='Value', hue='Variable', style='Variable', markers=False, dashes=False)
    palette = sns.color_palette()
    for i, (param, gt_value) in enumerate(ground_truth.items()):
        axes[1].hlines(y=gt_value, xmin=1, xmax=epoch, color=palette[i], linestyle='--')
    axes[1].set_title('Parameter Values (log) over Epochs with Ground Truth', fontsize=14)
    axes[1].set_xlabel('Epoch', fontsize=15)
    axes[1].set_ylabel('Value', fontsize=15)
    axes[1].set_yscale('log')
    axes[1].legend(prop={'size': 15})
    axes[1].tick_params(labelsize=14)
    axes[1].grid(True)

    plt.tight_layout()
    if save:
        plt.savefig(save_path + "training_log.png")

spatial_plot_surface(variable, coord, title='Variable', save_path='./', file_name='spatial_surface.png', grid_resolution=1000, method='CloughTocher', cmap='viridis', show=False)

Plots a smooth surface for spatial data.

The function interpolates scattered data onto a regular grid and generates a smooth surface plot. The resulting plot is saved as a PNG file.

Parameters:

Name	Type	Description	Default
variable	array - like	Values to plot, corresponding to the coordinates.	required
coord	array - like	Coordinates of shape (n, 2) where each row represents a point (x, y).	required
title	str	Title of the plot. Defaults to "Variable".	'Variable'
save_path	str	Directory to save the plot. Defaults to "./".	'./'
file_name	str	Name of the saved plot file. Defaults to "spatial_surface.png".	'spatial_surface.png'
grid_resolution	int	Resolution of the interpolation grid. Higher values produce finer surfaces. Defaults to 100.	1000
cmap	str	Colormap to use for the plot. Defaults to 'viridis'.	'viridis'

Raises:

Type	Description
ValueError	If the lengths of variable and coord do not match.

Returns:

Name	Type	Description
None	None	The function saves the plot to a file and does not return any value.

Example

import numpy as np coord = np.random.uniform(0, 10, (50, 2)) variable = np.sin(coord[:, 0]) + np.cos(coord[:, 1]) spatial_plot_surface(variable, coord, title="Example Plot", grid_resolution=200)

Source code in geospaNN/visualize.py

def spatial_plot_surface(variable: np.array,
                         coord: np.array,
                         title: Optional[str] = "Variable",
                         save_path: Optional[str] = "./",
                         file_name: Optional[str] = "spatial_surface.png",
                         grid_resolution: Optional[int] = 1000,
                         method: Optional[str] = "CloughTocher",
                         cmap: Optional[str] = 'viridis',
                         show: Optional[bool] = False
                         ) -> None:
    """
    Plots a smooth surface for spatial data.

    The function interpolates scattered data onto a regular grid and generates
    a smooth surface plot. The resulting plot is saved as a PNG file.

    Parameters:
        variable (array-like): Values to plot, corresponding to the coordinates.
        coord (array-like): Coordinates of shape (n, 2) where each row represents a point (x, y).
        title (str, optional): Title of the plot. Defaults to "Variable".
        save_path (str, optional): Directory to save the plot. Defaults to "./".
        file_name (str, optional): Name of the saved plot file. Defaults to "spatial_surface.png".
        grid_resolution (int, optional): Resolution of the interpolation grid.
            Higher values produce finer surfaces. Defaults to 100.
        cmap (str, optional): Colormap to use for the plot. Defaults to 'viridis'.

    Raises:
        ValueError: If the lengths of `variable` and `coord` do not match.

    Returns:
        None: The function saves the plot to a file and does not return any value.

    Example:
        >>> import numpy as np
        >>> coord = np.random.uniform(0, 10, (50, 2))
        >>> variable = np.sin(coord[:, 0]) + np.cos(coord[:, 1])
        >>> spatial_plot_surface(variable, coord, title="Example Plot", grid_resolution=200)
    """
    if len(variable) != len(coord):
        raise ValueError("Length of 'variable' and 'coord' must match.")

    # Create grid for interpolation
    xi = np.linspace(coord[:, 0].min(), coord[:, 0].max(), grid_resolution)
    yi = np.linspace(coord[:, 1].min(), coord[:, 1].max(), grid_resolution)
    X, Y = np.meshgrid(xi, yi)

    # Interpolate data onto the grid

    if method == "CloughTocher":
        interp = CloughTocher2DInterpolator(list(zip(coord[:, 0], coord[:, 1])), variable)
        Z = interp(X, Y)
    elif method in ['linear', 'nearest', 'cubic']:
        Z = griddata(coord, variable, (X, Y), method=method)
    else:
        warnings.warn("No interpolation method provided, use cubic spline as default!")
        Z = griddata(coord, variable, (X, Y), method="cubic")

    # Plot the interpolated surface
    plt.clf()
    fig = plt.figure(figsize=(8, 6))
    surface = plt.pcolormesh(X, Y, Z, cmap=cmap, shading='auto')
    plt.colorbar(surface, label='Variable')
    plt.title(title)
    plt.xlabel('Coord_X')
    plt.ylabel('Coord_Y')
    plt.savefig(f"{save_path}/{file_name}")
    if show:
        plt.show()

    return fig