torch-geopooling 1.2.0__cp311-cp311-manylinux_2_28_aarch64.whl

torch_geopooling/nn/embedding.py

@@ -0,0 +1,102 @@
+# Copyright (C) 2024, Yakau Bubnou
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from typing import Union, Tuple, cast
+
+import torch
+from torch import Tensor, nn
+
+from torch_geopooling import functional as F
+from torch_geopooling.tiling import Exterior, ExteriorTuple
+
+
+__all__ = [
+    "Embedding2d",
+]
+
+
+_Exterior = Union[Exterior, ExteriorTuple]
+
+
+class Embedding2d(nn.Module):
+    """
+    Retrieves spatial embeddings from a fixed-size lookup table based on 2D coordinates.
+
+    This module accepts a tensor of (x, y) coordinates and retrieves the corresponding
+    spatial embeddings from a provided embedding matrix. The embeddings are selected
+    based on the input coordinates, with an optional padding to include neighboring cells.
+
+    Args:
+        manifold: The size of the 2-dimensional embedding in a form (W, H, N), where
+            W is a width, H is a height, and N is a feature dimension of the embedding.
+        padding: The size of the neighborhood to query. Default is 0, meaning only the embedding
+            for the exact input coordinate is retrieved.
+        exterior: The geometric boundary of the learning space, specified as a tuple (X, Y, W, H),
+            where X and Y represent the origin, and W and H represent the width and height of the
+            space, respectively.
+        reflection: When true, kernel is wrapped around the exterior space, otherwise kernel is
+            squeezed into borders.
+
+    Shape:
+        - Input: :math:`(*, 2)`, where 2 comprises x and y coordinates.
+        - Output: :math:`(*, X_{out}, Y_{out}, N)`, where * is the input shape, \
+            :math:`N = \\text{manifold[2]}`, and
+
+            :math:`X_{out} = \\text{padding}[0] \\times 2 + 1`
+
+            :math:`Y_{out} = \\text{padding}[1] \\times 2 + 1`
+
+    Examples:
+
+    >>> # Create an embedding of EPSG:4326 rectangle into 1024x1024 embedding
+    >>> # with 3 features in each cell.
+    >>> embedding = nn.Embedding2d(
+    ...     (1024, 1024, 3),
+    ...     exterior=(-180.0, -90.0, 360.0, 180.0),
+    ...     padding=(2, 2),
+    ... )
+    >>> input = torch.rand((100, 2), dtype=torch.float64) * 60.0
+    >>> output = embedding(input)
+    """
+
+    def __init__(
+        self,
+        manifold: Tuple[int, int, int],
+        exterior: _Exterior,
+        padding: Tuple[int, int] = (0, 0),
+        reflection: bool = True,
+    ) -> None:
+        super().__init__()
+        self.manifold = manifold
+        self.exterior = cast(ExteriorTuple, tuple(map(float, exterior)))
+        self.padding = padding
+        self.reflection = reflection
+
+        self.weight = nn.Parameter(torch.empty(manifold, dtype=torch.float64))
+        nn.init.zeros_(self.weight)
+
+    def extra_repr(self) -> str:
+        return "{manifold}, exterior={exterior}, padding={padding}, reflection={reflection}".format(
+            **self.__dict__
+        )
+
+    def forward(self, input: Tensor) -> Tensor:
+        return F.embedding2d(
+            input,
+            self.weight,
+            exterior=self.exterior,
+            padding=self.padding,
+            reflection=self.reflection,
+        )

torch_geopooling/nn/embedding_test.py

@@ -0,0 +1,48 @@
+# Copyright (C) 2024, Yakau Bubnou
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import pytest
+import torch
+from torch import nn
+from torch.optim import SGD
+
+from torch_geopooling.nn.embedding import Embedding2d
+
+
+def test_embedding2d_optimize() -> None:
+    embedding = Embedding2d(
+        (2, 2, 1),
+        padding=(0, 0),
+        exterior=(-180.0, -90.0, 360.0, 180.0),
+    )
+
+    x_true = torch.tensor(
+        [[90.0, 45.0], [90.0, -45.0], [-90.0, -45.0], [-90.0, 45.0]], dtype=torch.float64
+    )
+    y_true = torch.tensor([[10.0], [20.0], [30.0], [40.0]], dtype=torch.float64)
+
+    optim = SGD(embedding.parameters(), lr=0.1)
+    loss_fn = nn.L1Loss()
+
+    for i in range(10000):
+        optim.zero_grad()
+
+        y_pred = embedding(x_true)
+        loss = loss_fn(y_pred[:, 0, 0, :], y_true)
+        loss.backward()
+
+        optim.step()
+
+    assert pytest.approx(0.0, abs=1e-1) == loss.detach().item()

torch_geopooling/nn/pooling.py

@@ -0,0 +1,396 @@
+# Copyright (C) 2024, Yakau Bubnou
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from typing import Optional, Union
+
+import torch
+from shapely.geometry import Polygon
+from torch import Tensor, nn
+
+from torch_geopooling import functional as F
+from torch_geopooling.tiling import Exterior, ExteriorTuple, regular_tiling
+
+__all__ = [
+    "AdaptiveAvgQuadPool2d",
+    "AdaptiveQuadPool2d",
+    "AdaptiveMaxQuadPool2d",
+    "AvgQuadPool2d",
+    "MaxQuadPool2d",
+    "QuadPool2d",
+]
+
+
+_Exterior = Union[Exterior, ExteriorTuple]
+
+
+_exterior_doc = """
+    Note:
+        Input coordinates must be within a specified exterior geometry (including boundaries).
+        For input coordinates outsize of the specified exterior, module throws an exception.
+"""
+
+
+_terminal_group_doc = """
+    Note:
+        A **terminal group** refers to a collection of terminal nodes within the quadtree that
+        share the same parent tile.
+"""
+
+
+class _AdaptiveQuadPool(nn.Module):
+    __doc__ = f"""
+    Args:
+        feature_dim: Size of each feature vector.
+        exterior: Geometrical boundary of the learning space in (X, Y, W, H) format.
+        max_terminal_nodes: Optional maximum number of terminal nodes in a quadtree. Once a
+            maximum is reached, internal nodes are no longer sub-divided and tree stops growing.
+        max_depth: Maximum depth of the quadtree. Default: 17.
+        capacity: Maximum number of inputs, after which a quadtree's node is subdivided and
+            depth of the tree grows. Default: 1.
+        precision: Optional rounding of the input coordinates. Default: 7.
+
+    Shape:
+        - Input: :math:`(*, 2)`, where 2 comprises longitude and latitude coordinates.
+        - Output: :math:`(*, H)`, where * is the input shape and :math:`H = \\text{{feature_dim}}`.
+
+    {_exterior_doc}
+    {_terminal_group_doc}
+    """
+
+    def __init__(
+        self,
+        feature_dim: int,
+        exterior: _Exterior,
+        max_terminal_nodes: Optional[int] = None,
+        max_depth: int = 17,
+        capacity: int = 1,
+        precision: Optional[int] = 7,
+    ) -> None:
+        super().__init__()
+        self.feature_dim = feature_dim
+        self.exterior = tuple(map(float, exterior))
+        self.max_terminal_nodes = max_terminal_nodes
+        self.max_depth = max_depth
+        self.capacity = capacity
+        self.precision = precision
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self) -> None:
+        # The weight for adaptive operation should be sparse, since training operation
+        # results in a dynamic change of the underlying quadtree.
+        weight_size = (
+            self.max_depth + 1,
+            1 << self.max_depth,
+            1 << self.max_depth,
+            self.feature_dim,
+        )
+        self.weight = nn.Parameter(torch.sparse_coo_tensor(size=weight_size, dtype=torch.float64))
+
+    @property
+    def tiles(self) -> torch.Tensor:
+        """Return tiles of the quadtree."""
+        return self.weight.coalesce().detach().indices().t()[:, :-1]
+
+    def extra_repr(self) -> str:
+        return (
+            "{feature_dim}, "
+            "exterior={exterior}, capacity={capacity}, max_depth={max_depth}, "
+            "precision={precision}".format(**self.__dict__)
+        )
+
+
+class AdaptiveQuadPool2d(_AdaptiveQuadPool):
+    __doc__ = f"""Adaptive lookup index over quadtree decomposition of input 2D coordinates.
+
+    This module constructs an internal lookup quadtree to organize closely situated 2D points.
+    Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
+    an input coordinate, the module retrieves the corresponding terminal node and returns its
+    associated weight.
+
+    {_AdaptiveQuadPool.__doc__}
+
+    Examples:
+
+    >>> # Feature vectors of size 4 over a 2d space.
+    >>> pool = nn.AdaptiveQuadPool2d(4, (-10, -5, 20, 10))
+    >>> # Grow tree up to 4-th level and sub-divides a node after 8 coordinates in a quad.
+    >>> pool = nn.AdaptiveQuadPool2d(4, (-10, -5, 20, 10), max_depth=4, capacity=8)
+    >>> # Create 2D coordinates and query associated weights.
+    >>> input = torch.rand((1024, 2), dtype=torch.float64) * 10 - 5
+    >>> output = pool(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        result = F.adaptive_quad_pool2d(
+            self.weight,
+            input,
+            self.exterior,
+            training=self.training,
+            max_terminal_nodes=self.max_terminal_nodes,
+            max_depth=self.max_depth,
+            capacity=self.capacity,
+            precision=self.precision,
+        )
+        if self.training:
+            self.weight.data = result.weight
+        return result.values
+
+
+class AdaptiveMaxQuadPool2d(_AdaptiveQuadPool):
+    __doc__ = f"""Adaptive maximum pooling over quadtree decomposition of input 2D coordinates.
+
+    This module constructs an internal lookup quadtree to organize closely situated 2D points.
+    Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
+    an input coordinate, the module retrieves a **terminal group** of nodes and calculates the
+    maximum value for each ``feature_dim``.
+
+    {_AdaptiveQuadPool.__doc__}
+
+    Examples:
+
+    >>> pool = nn.AdaptiveMaxQuadPool2d(3, (-10, -5, 20, 10), max_depth=5)
+    >>> # Create 2D coordinates and feature vector associated with them.
+    >>> input = torch.rand((2048, 2), dtype=torch.float64) * 10 - 5
+    >>> output = pool(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        result = F.adaptive_max_quad_pool2d(
+            self.weight,
+            input,
+            self.exterior,
+            training=self.training,
+            max_terminal_nodes=self.max_terminal_nodes,
+            max_depth=self.max_depth,
+            capacity=self.capacity,
+            precision=self.precision,
+        )
+        if self.training:
+            self.weight.data = result.weight
+        return result.values
+
+
+class AdaptiveAvgQuadPool2d(_AdaptiveQuadPool):
+    __doc__ = f"""Adaptive average pooling over quadtree decomposition of input 2D coordinates.
+
+    This module constructs an internal lookup quadtree to organize closely situated 2D points.
+    Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
+    an input coordinate, the module retrieves a **terminal group** of nodes and calculates an
+    average value for each ``feature_dim``.
+
+    {_AdaptiveQuadPool.__doc__}
+
+    Examples:
+
+    >>> # Create pool with 7 features.
+    >>> pool = nn.AdaptiveAvgQuadPool2d(7, (0, 0, 1, 1), max_depth=12)
+    >>> input = torch.rand((2048, 2), dtype=torch.float64)
+    >>> output = pool(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        result = F.adaptive_avg_quad_pool2d(
+            self.weight,
+            input,
+            self.exterior,
+            training=self.training,
+            max_terminal_nodes=self.max_terminal_nodes,
+            max_depth=self.max_depth,
+            capacity=self.capacity,
+            precision=self.precision,
+        )
+        if self.training:
+            self.weight.data = result.weight
+        return result.values
+
+
+class _QuadPool(nn.Module):
+    __doc__ = f"""
+    Args:
+        feature_dim: Size of each feature vector.
+        polygon: Polygon that resembles boundary for the terminal nodes of a quadtree.
+        exterior: Geometrical boundary of the learning space in (X, Y, W, H) format.
+        max_terminal_nodes: Optional maximum number of terminal nodes in a quadtree. Once a
+            maximum is reached, internal nodes are no longer sub-divided and tree stops growing.
+        max_depth: Maximum depth of the quadtree. Default: 17.
+        precision: Optional rounding of the input coordinates. Default: 7.
+
+    Shape:
+        - Input: :math:`(*, 2)`, where 2 comprises longitude and latitude coordinates.
+        - Output: :math:`(*, H)`, where * is the input shape and :math:`H = \\text{{feature_dim}}`.
+
+    {_exterior_doc}
+    {_terminal_group_doc}
+
+    Note:
+        All terminal nodes that have an intersection with the specified polygon boundary are
+        included into the quadtree.
+    """
+
+    def __init__(
+        self,
+        feature_dim: int,
+        polygon: Polygon,
+        exterior: _Exterior,
+        max_terminal_nodes: Optional[int] = None,
+        max_depth: int = 17,
+        precision: Optional[int] = 7,
+    ) -> None:
+        super().__init__()
+        self.feature_dim = feature_dim
+        self.polygon = polygon
+        self.exterior = tuple(map(float, exterior))
+        self.max_terminal_nodes = max_terminal_nodes
+        self.max_depth = max_depth
+        self.precision = precision
+
+        # Generate regular tiling for the provided polygon and build from those
+        # tiles a quadtree from terminal nodes all way up to the root node.
+        tiles_iter = regular_tiling(
+            polygon, Exterior.from_tuple(exterior), z=max_depth, internal=True
+        )
+        tiles = torch.tensor(list(tiles_iter), dtype=torch.int64)
+
+        self.register_buffer("tiles", tiles)
+        self.tiles: Tensor
+
+        self.initialize_parameters()
+        self.reset_parameters()
+
+    def initialize_parameters(self) -> None:
+        weight_size = [self.tiles.size(0), self.feature_dim]
+        self.weight = nn.Parameter(torch.empty(weight_size, dtype=torch.float64))
+
+    def reset_parameters(self) -> None:
+        nn.init.uniform_(self.weight)
+
+    def extra_repr(self) -> str:
+        return (
+            "{feature_dim}, exterior={exterior}, max_depth={max_depth}, "
+            "precision={precision}".format(**self.__dict__)
+        )
+
+
+class QuadPool2d(_QuadPool):
+    __doc__ = f"""Lookup index over quadtree decomposition of input 2D coordinates.
+
+    This module constructs an internal lookup tree to organize closely situated 2D points using
+    a specified polygon and exterior, where polygon is treated as a *boundary* of terminal
+    nodes of a quadtree.
+
+    Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
+    an input coordinate, the module retrieves the corresponding terminal node and returns its
+    associated weight.
+
+    {_QuadPool.__doc__}
+
+    Examples:
+
+    >>> from shapely.geometry import Polygon
+    >>> # Create a pool for squared exterior 100x100 and use only a portion of that
+    >>> # exterior isolated by a square 10x10.
+    >>> poly = Polygon([(0, 0), (10, 0), (10, 10), (0, 10)])
+    >>> pool = nn.QuadPool2d(5, poly, exterior=(0, 0, 100, 100))
+    >>> input = torch.rand((2048, 2), dtype=torch.float64)
+    >>> output = pool(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        result = F.quad_pool2d(
+            self.tiles,
+            self.weight,
+            input,
+            self.exterior,
+            # This is not a mistake, since we already know the shape of the
+            # quadtree, there is no need to learn it.
+            training=False,
+            max_terminal_nodes=self.max_terminal_nodes,
+            max_depth=self.max_depth,
+            precision=self.precision,
+        )
+        return result.values
+
+
+class MaxQuadPool2d(_QuadPool):
+    __doc__ = f"""Maximum pooling over quadtree decomposition of input 2D coordinates.
+
+    This module constructs an internal lookup tree to organize closely situated 2D points using
+    a specified polygon and exterior, where polygon is treated as a *boundary* of terminal nodes
+    of a quadtree.
+
+    Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
+    an input coordinate, the module retrieves a **terminal group** of nodes and calculates the
+    maximum value for each ``feature_dim``.
+
+    {_QuadPool.__doc__}
+
+    Examples:
+
+    >>> from shapely.geometry import Polygon
+    >>> poly = Polygon([(0, 0), (10, 0), (10, 10), (0, 10)])
+    >>> pool = nn.MaxQuadPool2d(3, poly, exterior=(0, 0, 100, 100))
+    >>> input = torch.rand((2048, 2), dtype=torch.float64)
+    >>> output = pool(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        result = F.max_quad_pool2d(
+            self.tiles,
+            self.weight,
+            input,
+            self.exterior,
+            training=False,
+            max_terminal_nodes=self.max_terminal_nodes,
+            max_depth=self.max_depth,
+            precision=self.precision,
+        )
+        return result.values
+
+
+class AvgQuadPool2d(_QuadPool):
+    __doc__ = f"""Average pooling over quadtree decomposition of input 2D coordinates.
+
+    This module constructs an internal lookup tree to organize closely situated 2D points using
+    a specified polygon and exterior, where polygon is treated as a *boundary* of terminal
+    nodes of a quadtree.
+
+    Each terminal node in the resulting quadtree is paired with a weight. Thus, when providing
+    an input coordinate, the module retrieves a **terminal group** of nodes and calculates an
+    average value for each ``feature_dim``.
+
+    {_QuadPool.__doc__}
+
+    Examples:
+
+    >>> from shapely.geometry import Polygon
+    >>> poly = Polygon([(0, 0), (10, 0), (10, 10), (0, 10)])
+    >>> pool = nn.AvgQuadPool2d(4, poly, exterior=(0, 0, 100, 100))
+    >>> input = torch.rand((2048, 2), dtype=torch.float64)
+    >>> output = pool(input)
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        result = F.avg_quad_pool2d(
+            self.tiles,
+            self.weight,
+            input,
+            self.exterior,
+            training=False,
+            max_terminal_nodes=self.max_terminal_nodes,
+            max_depth=self.max_depth,
+            precision=self.precision,
+        )
+        return result.values

torch_geopooling/nn/pooling_test.py

@@ -0,0 +1,158 @@
+# Copyright (C) 2024, Yakau Bubnou
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from typing import Type
+
+import pytest
+import torch
+from shapely.geometry import Polygon
+from torch import nn
+from torch.optim import SGD
+from torch.nn import L1Loss
+
+from torch_geopooling.nn.pooling import (
+    AdaptiveAvgQuadPool2d,
+    AdaptiveMaxQuadPool2d,
+    AdaptiveQuadPool2d,
+    AvgQuadPool2d,
+    MaxQuadPool2d,
+    QuadPool2d,
+)
+
+
+@pytest.mark.parametrize(
+    "module_class",
+    [
+        AdaptiveQuadPool2d,
+        AdaptiveMaxQuadPool2d,
+        AdaptiveAvgQuadPool2d,
+    ],
+    ids=["id", "max", "avg"],
+)
+def test_adaptive_quad_pool2d_gradient(module_class: Type[nn.Module]) -> None:
+    pool = module_class(5, (-180, -90, 360, 180))
+
+    input = torch.rand((100, 2), dtype=torch.float64) * 90
+    y = pool(input)
+
+    assert pool.weight.grad is None
+
+    loss_fn = L1Loss()
+    loss = loss_fn(y, torch.ones_like(y))
+    loss.backward()
+
+    assert pool.weight.grad is not None
+    assert pool.weight.grad.sum().item() == pytest.approx(-1)
+
+
+def test_adaptive_quad_pool2d_optimize() -> None:
+    pool = AdaptiveQuadPool2d(1, (-180, -90, 360, 180), max_depth=1)
+
+    # Input coordinates are simply centers of the level-1 quads.
+    x_true = torch.tensor(
+        [[90.0, 45.0], [90.0, -45.0], [-90.0, -45.0], [-90.0, 45.0]], dtype=torch.float64
+    )
+    y_true = torch.tensor([[10.0], [20.0], [30.0], [40.0]], dtype=torch.float64)
+    y_tile = [[1, 1, 1], [1, 1, 0], [1, 0, 0], [1, 0, 1]]
+
+    optim = SGD(pool.parameters(), lr=0.01)
+    loss_fn = nn.L1Loss()
+
+    for i in range(20000):
+        optim.zero_grad()
+
+        y_pred = pool(x_true)
+        loss = loss_fn(y_pred, y_true)
+        loss.backward()
+
+        optim.step()
+
+    # Ensure that model converged with a small loss.
+    assert pytest.approx(0.0, abs=1e-1) == loss.detach().item()
+
+    # Ensure that weights that pooling operation learned are the same as in the
+    # target matrix (y_true).
+    weight = pool.weight.to_dense()
+
+    for i, tile in enumerate(y_tile):
+        z, x, y = tile
+        expect_weight = y_true[i].item()
+        actual_weight = weight[z, x, y].detach().item()
+
+        assert pytest.approx(expect_weight, abs=1e-1) == actual_weight, f"tile {tile} is wrong"
+
+
+@pytest.mark.parametrize(
+    "module_class",
+    [
+        QuadPool2d,
+        MaxQuadPool2d,
+        AvgQuadPool2d,
+    ],
+    ids=["id", "max", "avg"],
+)
+def test_quad_pool2d_gradient(module_class: Type[nn.Module]) -> None:
+    poly = Polygon([(0.0, 0.0), (1.0, 0.0), (1.0, 1.1), (0.0, 1.0)])
+    exterior = (0.0, 0.0, 1.0, 1.0)
+
+    pool = module_class(4, poly, exterior, max_depth=5)
+    assert pool.weight.size() == torch.Size([pool.tiles.size(0), 4])
+
+    input = torch.rand((100, 2), dtype=torch.float64)
+    y = pool(input)
+
+    assert pool.weight.grad is None
+
+    loss_fn = L1Loss()
+    loss = loss_fn(y, torch.ones_like(y))
+    loss.backward()
+
+    assert pool.weight.grad is not None
+    assert pool.weight.grad.sum().item() == pytest.approx(-1)
+
+
+def test_quad_pool2d_optimize() -> None:
+    poly = Polygon([(-180, -90), (-180, 90), (180, 90), (180, -90)])
+    pool = QuadPool2d(1, poly, (-180, -90, 360, 180), max_depth=1)
+
+    x_true = torch.tensor(
+        [[90.0, 45.0], [90.0, -45.0], [-90.0, -45.0], [-90.0, 45.0]], dtype=torch.float64
+    )
+    y_true = torch.tensor([[10.0], [20.0], [30.0], [40.0]], dtype=torch.float64)
+    y_tile = [(1, 1, 1), (1, 1, 0), (1, 0, 0), (1, 0, 1)]
+
+    optim = SGD(pool.parameters(), lr=0.01)
+    loss_fn = nn.L1Loss()
+
+    for i in range(20000):
+        optim.zero_grad()
+
+        y_pred = pool(x_true)
+        loss = loss_fn(y_pred, y_true)
+        loss.backward()
+
+        optim.step()
+
+    # Ensure that model converged with a small loss.
+    assert pytest.approx(0.0, abs=1e-1) == loss.detach().item()
+
+    actual_tiles = {}
+    for i in range(pool.tiles.size(0)):
+        tile = tuple(pool.tiles[i].detach().tolist())
+        actual_tiles[tile] = pool.weight[i, 0].detach().item()
+
+    for tile, expect_weight in zip(y_tile, y_true[:, 0].tolist()):
+        actual_weight = actual_tiles[tile]
+        assert pytest.approx(expect_weight, abs=1e-1) == actual_weight, f"tile {tile} is wrong"