Skip to content

Dataset

embedders.geovex.dataset

HexagonalDataset.

This dataset is used to train a hexagonal encoder model. As defined in GeoVex paper[1].

References

[1] https://openreview.net/forum?id=7bvWopYY1H

HexagonalDataset 

HexagonalDataset(
    data: pd.DataFrame, neighbourhood: H3Neighbourhood, neighbor_k_ring: int = 6
)

Bases: Dataset['torch.Tensor'], Generic[T]

Dataset for the hexagonal encoder model.

It works by returning a 3d tensor of hexagonal regions. The tensor is a cube with the target hexagonal region in the center, and the rings of neighbors around surrounding it.

PARAMETER DESCRIPTION
data

Data to use for training. Raw counts of features in regions.

TYPE: DataFrame

neighbourhood

H3Neighbourhood to use for training. It has to be initialized with the same data as the data argument.

TYPE: H3Neighbourhood

neighbor_k_ring

The hexagonal rings of neighbors to include in the tensor. Defaults to 6.

TYPE: int DEFAULT: 6

Source code in srai/embedders/geovex/dataset.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def __init__(
    self,
    data: pd.DataFrame,
    neighbourhood: H3Neighbourhood,
    neighbor_k_ring: int = 6,
):
    """
    Initialize the HexagonalDataset.

    Args:
        data (pd.DataFrame): Data to use for training. Raw counts of features in regions.
        neighbourhood (H3Neighbourhood): H3Neighbourhood to use for training.
            It has to be initialized with the same data as the data argument.
        neighbor_k_ring (int, optional): The hexagonal rings of neighbors to include
            in the tensor. Defaults to 6.
    """
    import_optional_dependencies(dependency_group="torch", modules=["torch"])
    import torch

    self._assert_k_ring_correct(neighbor_k_ring)
    self._assert_h3_neighbourhood(neighbourhood)
    # store the desired k
    self._k: int = neighbor_k_ring
    # number of columns in the dataset
    self._N: int = data.shape[1]
    # store the list of valid h3 indices (have all the neighbors in the dataset)
    self._valid_cells: list[CellInfo] = []
    # store the data as a torch tensor
    self._data_torch = torch.Tensor(data.to_numpy(dtype=np.float32))
    # iterate over the data and build the valid h3 indices
    self._invalid_cells, self._valid_cells = self._seperate_valid_invalid_cells(
        data, neighbourhood, neighbor_k_ring, set(data.index)
    )

__getitem__ 

__getitem__(index: Any) -> torch.Tensor

Return a single item from the dataset.

PARAMETER DESCRIPTION
index

The index of dataset item to return

TYPE: Any

RETURNS DESCRIPTION
HexagonalDatasetItem

The dataset item

TYPE: Tensor

Source code in srai/embedders/geovex/dataset.py 
124
125
126
127
128
129
130
131
132
133
134
135
def __getitem__(self, index: Any) -> "torch.Tensor":
    """
    Return a single item from the dataset.

    Args:
        index (Any): The index of dataset item to return

    Returns:
        HexagonalDatasetItem: The dataset item
    """
    _, target_idx, neighbors_idxs = self._valid_cells[index]
    return self._build_tensor(target_idx, neighbors_idxs)

__len__ 

__len__() -> int

Returns the number of valid h3 indices in the dataset.

RETURNS DESCRIPTION
int

Number of valid h3 indices in the dataset.

TYPE: int

Source code in srai/embedders/geovex/dataset.py 
115
116
117
118
119
120
121
122
def __len__(self) -> int:
    """
    Returns the number of valid h3 indices in the dataset.

    Returns:
        int: Number of valid h3 indices in the dataset.
    """
    return len(self._valid_cells)

get_invalid_cells 

get_invalid_cells() -> list[str]

Returns the list of invalid h3 indices in the dataset.

RETURNS DESCRIPTION
list[str]

List[str]: List of invalid h3 indices in the dataset.
Source code in srai/embedders/geovex/dataset.py 
190
191
192
193
194
195
196
197
def get_invalid_cells(self) -> list[str]:
    """
    Returns the list of invalid h3 indices in the dataset.

    Returns:
        List[str]: List of invalid h3 indices in the dataset.
    """
    return list(self._invalid_cells)

get_valid_cells 

get_valid_cells() -> list[str]

Returns the list of valid h3 indices in the dataset.

RETURNS DESCRIPTION
list[str]

List[str]: List of valid h3 indices in the dataset.
Source code in srai/embedders/geovex/dataset.py 
181
182
183
184
185
186
187
188
def get_valid_cells(self) -> list[str]:
    """
    Returns the list of valid h3 indices in the dataset.

    Returns:
        List[str]: List of valid h3 indices in the dataset.
    """
    return [h3_index for h3_index, _, _ in self._valid_cells]