From ae3c960217c8f8349f9ab30f31be4b73b251b1d6 Mon Sep 17 00:00:00 2001 From: Zain-Mahmoud Date: Wed, 12 Aug 2026 01:17:03 +0300 Subject: [PATCH 01/10] add tests --- tests/test_inference.py | 43 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 tests/test_inference.py diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..3367b15 --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,43 @@ +""" +Tests inference interface by instantiating mock models +and implementing the `generate` function +""" +from terrain_diffusion.inference import TerrainModel +import numpy as np +import pytest + +class MockTerrainModel(TerrainModel): + """ + Mock model implementing the TerrainModel abstract class + """ + def generate(self, patch: np.ndarray) -> np.ndarray: + """ + Generate and output patch (C, H, W) from the input patch + by multiplying by 2 + """ + return patch * 2 + + @classmethod + def load_model(cls, model_path): + print(f"loaded model: {model_path}") + return MockTerrainModel() + + def __call__(self, patch: np.ndarray) -> np.ndarray: + return self.generate(patch) + + +class TestTerrainModel: + @pytest.fixture + def initial_model(self) -> MockTerrainModel: + return MockTerrainModel() + + def test_load(self): + assert isinstance(MockTerrainModel.load_model("my_model"), TerrainModel) + + def test_generate(self, initial_model): + input_patch = np.ones((3, 5, 10)) # 10x5 image with 3 channels + + actual = initial_model(input_patch) + expected = 2 * input_patch + + assert np.array_equal(actual, expected), "patches are not equal" From 0de35c4aed5b3dead087aed3fde8b49cae6b12d1 Mon Sep 17 00:00:00 2001 From: Zain-Mahmoud Date: Wed, 12 Aug 2026 01:17:20 +0300 Subject: [PATCH 02/10] create model interface --- src/terrain_diffusion/inference.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/terrain_diffusion/inference.py b/src/terrain_diffusion/inference.py index 948f8d1..5c2bbb7 100644 --- a/src/terrain_diffusion/inference.py +++ b/src/terrain_diffusion/inference.py @@ -16,3 +16,26 @@ - It loads weights from the external model weights download. - It runs on the GPU compute node. """ +import numpy as np +from abc import ABC, abstractmethod + +class TerrainModel(ABC): + + @abstractmethod + def generate(self, patch: np.ndarray) -> np.ndarray: + """ + Generate an output patch (C, H, W) from an input patch (C, H, W) using the model + """ + raise NotImplemented + + @classmethod + @abstractmethod + def load_model(cls, model_path: str) -> TerrainModel: + """ + Load a model stored in model_path + """ + raise NotImplemented + + def __call__(self, patch: np.ndarray) -> np.ndarray: + return self.generate(patch) + From dadec335fbb812232107c8bc6ef9f2f74b5805e9 Mon Sep 17 00:00:00 2001 From: Zain-Mahmoud Date: Wed, 12 Aug 2026 01:18:50 +0300 Subject: [PATCH 03/10] update formatting --- src/terrain_diffusion/inference.py | 10 ++++++---- tests/test_inference.py | 4 +++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/src/terrain_diffusion/inference.py b/src/terrain_diffusion/inference.py index 5c2bbb7..16e153e 100644 --- a/src/terrain_diffusion/inference.py +++ b/src/terrain_diffusion/inference.py @@ -16,9 +16,11 @@ - It loads weights from the external model weights download. - It runs on the GPU compute node. """ -import numpy as np from abc import ABC, abstractmethod +import numpy as np + + class TerrainModel(ABC): @abstractmethod @@ -26,7 +28,7 @@ def generate(self, patch: np.ndarray) -> np.ndarray: """ Generate an output patch (C, H, W) from an input patch (C, H, W) using the model """ - raise NotImplemented + raise NotImplementedError @classmethod @abstractmethod @@ -34,8 +36,8 @@ def load_model(cls, model_path: str) -> TerrainModel: """ Load a model stored in model_path """ - raise NotImplemented - + raise NotImplementedError + def __call__(self, patch: np.ndarray) -> np.ndarray: return self.generate(patch) diff --git a/tests/test_inference.py b/tests/test_inference.py index 3367b15..a73ee8b 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -2,10 +2,12 @@ Tests inference interface by instantiating mock models and implementing the `generate` function """ -from terrain_diffusion.inference import TerrainModel import numpy as np import pytest +from terrain_diffusion.inference import TerrainModel + + class MockTerrainModel(TerrainModel): """ Mock model implementing the TerrainModel abstract class From 8c31061978177528eba332756515da14c5f36719 Mon Sep 17 00:00:00 2001 From: Zain-Mahmoud Date: Wed, 12 Aug 2026 01:26:07 +0300 Subject: [PATCH 04/10] fix formatting --- src/terrain_diffusion/inference.py | 3 +-- tests/test_inference.py | 4 +++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/terrain_diffusion/inference.py b/src/terrain_diffusion/inference.py index 16e153e..4820ab2 100644 --- a/src/terrain_diffusion/inference.py +++ b/src/terrain_diffusion/inference.py @@ -16,13 +16,13 @@ - It loads weights from the external model weights download. - It runs on the GPU compute node. """ + from abc import ABC, abstractmethod import numpy as np class TerrainModel(ABC): - @abstractmethod def generate(self, patch: np.ndarray) -> np.ndarray: """ @@ -40,4 +40,3 @@ def load_model(cls, model_path: str) -> TerrainModel: def __call__(self, patch: np.ndarray) -> np.ndarray: return self.generate(patch) - diff --git a/tests/test_inference.py b/tests/test_inference.py index a73ee8b..63c96ea 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -2,6 +2,7 @@ Tests inference interface by instantiating mock models and implementing the `generate` function """ + import numpy as np import pytest @@ -12,9 +13,10 @@ class MockTerrainModel(TerrainModel): """ Mock model implementing the TerrainModel abstract class """ + def generate(self, patch: np.ndarray) -> np.ndarray: """ - Generate and output patch (C, H, W) from the input patch + Generate an output patch (C, H, W) from the input patch by multiplying by 2 """ return patch * 2 From 1c26fcfd6e7d91ff49bc7b62cc75c328ac1212d9 Mon Sep 17 00:00:00 2001 From: Zain-Mahmoud Date: Wed, 12 Aug 2026 01:27:26 +0300 Subject: [PATCH 05/10] fix comment formatting --- tests/test_inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_inference.py b/tests/test_inference.py index 63c96ea..5ef4fde 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -39,7 +39,7 @@ def test_load(self): assert isinstance(MockTerrainModel.load_model("my_model"), TerrainModel) def test_generate(self, initial_model): - input_patch = np.ones((3, 5, 10)) # 10x5 image with 3 channels + input_patch = np.ones((3, 5, 10)) # 10x5 image with 3 channels actual = initial_model(input_patch) expected = 2 * input_patch From 27f73621db6d92d21ef3a016f477ddfd1c0fbc10 Mon Sep 17 00:00:00 2001 From: Zain-Mahmoud Date: Wed, 12 Aug 2026 01:29:25 +0300 Subject: [PATCH 06/10] update dependencies --- pyproject.toml | 4 +++- uv.lock | 55 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index fcc0f72..ea7b2a9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,9 @@ requires-python = ">=3.14" # Packages the project itself needs in order to run # TODO: add entries here as the project starts needing them, # then run `uv sync` to install and update uv.lock -dependencies = [] +dependencies = [ + "numpy>=2.5.2", +] # The commands the project installs. `uv run terrain-diffusion` runs main() in # src/terrain_diffusion/cli.py. diff --git a/uv.lock b/uv.lock index fe691f7..ed02485 100644 --- a/uv.lock +++ b/uv.lock @@ -20,6 +20,57 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "numpy" +version = "2.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/80/db0b4559e57ec36362bedbb05530a87fafbcb6067708c946967a41d449e7/numpy-2.5.2.tar.gz", hash = "sha256:d482d171c406ae88c5b19cad3b6a1c4c5209f886ab74bc44c2c865c23f52d860", size = 20773161, upload-time = "2026-08-09T13:48:27.962Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ac/f8/c3b222bf075b50afd8e949a07a15c4b312a4a84bd8102a332bcd953cbbb4/numpy-2.5.2-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:d787cf769c3baeb5f6235e778edb52c08dfa923789b5958f28e6450f96107cb1", size = 16885180, upload-time = "2026-08-09T13:46:03.939Z" }, + { url = "https://files.pythonhosted.org/packages/17/e1/2c1d4b1987795a92b5bbf7c24fe249ab96aa2573ab0d7604802c189d7b86/numpy-2.5.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:24b9dc2e3d84aa58523798805194e23e736f3f6ce2d1a5b92583ae734e6dbda8", size = 11907878, upload-time = "2026-08-09T13:46:07.045Z" }, + { url = "https://files.pythonhosted.org/packages/b9/ee/d08226fc858044355983a6e5b94f08ff6f3969e0a2b160a4a89f0ddb3445/numpy-2.5.2-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:9e9413326d726c2545bfa65d2c0876871e8d8386e77f992c1d426e180bbd4323", size = 5354922, upload-time = "2026-08-09T13:46:10.04Z" }, + { url = "https://files.pythonhosted.org/packages/94/f0/6d3d933056440ebbc5e6bad92065fc6c26a48a84a36b1208580e94eea76c/numpy-2.5.2-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:60e902ac295855348a5ca2ea4c89108989a9f5fddfad3dfc0a8f36b10358567e", size = 6679168, upload-time = "2026-08-09T13:46:12.275Z" }, + { url = "https://files.pythonhosted.org/packages/c4/3b/ecd49dd90033cceb2704d88ca905d4d7d89b0e8c739608754ffd325fa820/numpy-2.5.2-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:50e500dc868e9313530ce12ba470fe50ff3afe3d62993ed6eff652dacd555b65", size = 15624501, upload-time = "2026-08-09T13:46:15.322Z" }, + { url = "https://files.pythonhosted.org/packages/c7/99/461bd36dbdfac6c1c53efa370bd55a83227542d0d118f1677dbf1a3dacd5/numpy-2.5.2-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:318b9a4c845dbea06708a29c84ee429cc3065048db34cdb799047643492050ee", size = 16713701, upload-time = "2026-08-09T13:46:18.949Z" }, + { url = "https://files.pythonhosted.org/packages/f9/9c/2b251df9e8a5d647b62b0cbc1b90a91850c1cf4859ecb532fd0b4eacff6c/numpy-2.5.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:34c319e2963be042673fb46570501b2f06c41924e17e3563d58646b4380dfb68", size = 16986065, upload-time = "2026-08-09T13:46:23.006Z" }, + { url = "https://files.pythonhosted.org/packages/8f/25/20de43f53ff1390534a124475055a19f01fe10c920a0fd11b8e18d6d6052/numpy-2.5.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f06571a052127dc1b4e8b83029b4d1b20daa2b64a31cdd181fc6bc774e9000eb", size = 18470031, upload-time = "2026-08-09T13:46:27.102Z" }, + { url = "https://files.pythonhosted.org/packages/56/5e/0c577ca308d6da5eb79b546ba10bbe5b60148192194e2da060913b1de4f1/numpy-2.5.2-cp314-cp314-win32.whl", hash = "sha256:2cc779226e476d1e1f08c74068c419e60f41a9e0e069c92f6671d31d5c985e98", size = 6121028, upload-time = "2026-08-09T13:46:30.046Z" }, + { url = "https://files.pythonhosted.org/packages/15/5c/7bcbd5b11f94199073320410cddcbb80cee62415bfeb540874b265c2d922/numpy-2.5.2-cp314-cp314-win_amd64.whl", hash = "sha256:7587f53dfbd5edc0f7b87c6217b4c6d2d1f2ef9c3da70bc1315e7db5f8d7ec9d", size = 12597627, upload-time = "2026-08-09T13:46:32.886Z" }, + { url = "https://files.pythonhosted.org/packages/87/bc/4d0b06fba0da90ccc75af62823cb9dcedb6c9ea0cffa058cb2c9ee773a77/numpy-2.5.2-cp314-cp314-win_arm64.whl", hash = "sha256:3e4c367352d3747784248a227fbec218e193b56f7e6692e3b64fc805478ecfdf", size = 10680414, upload-time = "2026-08-09T13:46:36.036Z" }, + { url = "https://files.pythonhosted.org/packages/cd/17/f429aac9dc08833a0d0f188eba38c532a751b1a1f2ca6018a37b455cb321/numpy-2.5.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:b879fb674276e331513fb136b78dbc6bd3c848309e0d841cfd63be3896c4cfc1", size = 12026967, upload-time = "2026-08-09T13:46:39.084Z" }, + { url = "https://files.pythonhosted.org/packages/ca/9f/d0849de96a2a4ceaa16662f18ee13eaa9c0aa418269fdc8c4857c56b11da/numpy-2.5.2-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:fd0d703772bba096843785bd38371e31bb4a0c1151497ad5739d182114a73f7f", size = 5473874, upload-time = "2026-08-09T13:46:42.075Z" }, + { url = "https://files.pythonhosted.org/packages/89/3c/8df216d4a4a5422a3de045301cf7df8ea47286d76f5cb7160b0128ac26b7/numpy-2.5.2-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:3a2f061cebd9e3d23bdcfaaded5e2293a4c6a5b60fa42df85d410a725ce621bf", size = 6789276, upload-time = "2026-08-09T13:46:44.387Z" }, + { url = "https://files.pythonhosted.org/packages/e6/3a/20d7e9891c4ddfadd6ff8d95bf4b29f353d8e1770553de2099880551dfb9/numpy-2.5.2-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6df895598c0edcb41030126c89e0f353b07d93238116143b7405e937359736c4", size = 15659154, upload-time = "2026-08-09T13:46:47.538Z" }, + { url = "https://files.pythonhosted.org/packages/aa/d6/f3aa3d2688bf501b858835c6bd087ae9b51a56ae6fca8e2b0990abd177af/numpy-2.5.2-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1ab3d4a901f844ea836c3e80bf463c6a27d7f3c14e8e292fcf28d348b25b9bce", size = 16748909, upload-time = "2026-08-09T13:46:51.442Z" }, + { url = "https://files.pythonhosted.org/packages/7d/8f/1c5cae8d2baf86ab802ae97a00be55bc7e21ebc11b12bbc33376c5f05342/numpy-2.5.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:cebc2d6dbb605a7703d59751dea4bd6b0ab127a5a4338a6f432df1936fef8b26", size = 17027685, upload-time = "2026-08-09T13:46:55.095Z" }, + { url = "https://files.pythonhosted.org/packages/5c/27/71d3467404aedc1c24ce79610f91b52b0b0f466c43a701aa56fc75c145ab/numpy-2.5.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:eaca7ff36f0f52e2111ec71f169d8fd3e889e7ddc0d2592e0d703fd8d3ce8fac", size = 18501181, upload-time = "2026-08-09T13:46:59.09Z" }, + { url = "https://files.pythonhosted.org/packages/14/2f/42921d27c40aea7e077f4a423ae509fd9220b028cd787bafefd8ab2b3a5f/numpy-2.5.2-cp314-cp314t-win32.whl", hash = "sha256:ddf47472af2e4280d79bac82304f5e80150211f1b9e614b760061d5fdfbb6eba", size = 6271085, upload-time = "2026-08-09T13:47:01.903Z" }, + { url = "https://files.pythonhosted.org/packages/75/e6/bad5f5d56de9b1971bac959963dda276d35c40f1854475005434bbe08692/numpy-2.5.2-cp314-cp314t-win_amd64.whl", hash = "sha256:44ef9675d908e65f9953063837c3277730f3f4437615a4cdab67b366cabaf884", size = 12787971, upload-time = "2026-08-09T13:47:04.963Z" }, + { url = "https://files.pythonhosted.org/packages/df/05/f608795cb34391acd67e38d94a3c36abd8d8576293a3a80727d7595c372c/numpy-2.5.2-cp314-cp314t-win_arm64.whl", hash = "sha256:eaa088384c46f519dacb93b7ec483a6d6b19a4a2085ae4f25ab9b1c43d387d1e", size = 10750306, upload-time = "2026-08-09T13:47:07.976Z" }, + { url = "https://files.pythonhosted.org/packages/33/c6/28de0191c5f82b7d42a0a51390ba98587048aa93a39fafb05bdbe6e8d00c/numpy-2.5.2-cp315-cp315-macosx_10_15_x86_64.whl", hash = "sha256:078f9b027b478c9379b9677babbf0f8b8f1ecfada27636d7b9a93990c638739f", size = 16885274, upload-time = "2026-08-09T13:47:11.439Z" }, + { url = "https://files.pythonhosted.org/packages/dd/d1/973ca116000d244897e468ea1aff30b589e5022e3c8744b71706fe33bd57/numpy-2.5.2-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:50a68f4bacd8a2b33d8da3d2269d0d78500f86ea582e4786dc10f5ef2c2c6842", size = 11907846, upload-time = "2026-08-09T13:47:15.128Z" }, + { url = "https://files.pythonhosted.org/packages/78/d9/8c4b3937ef204cb2fd88d389ccd0f265a2ffb11f35a01d2064cf46714bd6/numpy-2.5.2-cp315-cp315-macosx_14_0_arm64.whl", hash = "sha256:e79aba74ffaf5f78a050d777c184cddf8fdffabab38acf5f3ef1fecbc17895d6", size = 5354892, upload-time = "2026-08-09T13:47:18.07Z" }, + { url = "https://files.pythonhosted.org/packages/74/9b/b6ee65ea2999fdb7023935e108e6fb776ee4082aa15f159acfa857e578c8/numpy-2.5.2-cp315-cp315-macosx_14_0_x86_64.whl", hash = "sha256:9a0731745a72a184490a582fb4af2533512bd071ace67785b5fdffc0ae58dce8", size = 6679309, upload-time = "2026-08-09T13:47:20.456Z" }, + { url = "https://files.pythonhosted.org/packages/43/f3/acb18d8b137a393c8e7803a8c994c9e64bde3930692a69d826993113a159/numpy-2.5.2-cp315-cp315-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4ec954036759bcee3aa484f8603bd9c14f3e776293b85578b8734c2d72777c69", size = 15625850, upload-time = "2026-08-09T13:47:24.365Z" }, + { url = "https://files.pythonhosted.org/packages/a9/bf/a8e9bb0db815a0e265b5744ebedd3af0bd5faad8604e5b50a1cd012f3c91/numpy-2.5.2-cp315-cp315-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dc649493697006bc90614a5f0bbc8cb3cb1866715c474e473694968d7e6b99ab", size = 16713664, upload-time = "2026-08-09T13:47:27.965Z" }, + { url = "https://files.pythonhosted.org/packages/0c/c3/6e913736b3dd6582344af32418b5fb9dab34282e8a8174ae1d54ceb0fc13/numpy-2.5.2-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:cf7de32f486e4ac9e2d93b810f9e9ac72a728dd46a32a0bb403222f27f653514", size = 16986749, upload-time = "2026-08-09T13:47:31.541Z" }, + { url = "https://files.pythonhosted.org/packages/80/09/7d3b23eff5c7428ef6c01e6f7052bb60d504c4d33e317b36b8959c24ad97/numpy-2.5.2-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:2ffa7bacab3e2ee1b19ed31766bb60bb380b68c23f051e199c5cc598afd68710", size = 18470495, upload-time = "2026-08-09T13:47:35.364Z" }, + { url = "https://files.pythonhosted.org/packages/a5/a4/68a321d825374f6eb677ffe8ef8c6b9a328304e6fd2e39d9530822776607/numpy-2.5.2-cp315-cp315-win32.whl", hash = "sha256:6b588cc8f902d6bff201c19fd00c43ab8545671e3554d014e12e14139e5e8617", size = 6120696, upload-time = "2026-08-09T13:47:38.561Z" }, + { url = "https://files.pythonhosted.org/packages/c8/23/deafbb1700f79fae9cd1e91220f133d124cc267de1b584da3fbf6db2f6cd/numpy-2.5.2-cp315-cp315-win_amd64.whl", hash = "sha256:07d4e89f3a9ab0a9ba24264ccdb642b3dd951b2281e8883a5481a4aa79cc31a7", size = 12597324, upload-time = "2026-08-09T13:47:41.401Z" }, + { url = "https://files.pythonhosted.org/packages/33/cd/3272ba105e3bbbdaeb11357eda31e7a6825ffe159e8171665660299a948f/numpy-2.5.2-cp315-cp315-win_arm64.whl", hash = "sha256:a610dc7e3c52edd39c2bc2375ff9c3fd59cb3ad00e4472d36f83bc1457145788", size = 10680466, upload-time = "2026-08-09T13:47:44.873Z" }, + { url = "https://files.pythonhosted.org/packages/0e/0e/58370637b1bb70a5c9ce2b43f4b521ccb224e36ccb76a6596b17ae4b447c/numpy-2.5.2-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:40f4d451aed46a8046a1aae41c4e55fb3612273df9c502480135e1501576a34b", size = 16993947, upload-time = "2026-08-09T13:47:48.97Z" }, + { url = "https://files.pythonhosted.org/packages/10/93/2abcb807712b289d6d60fe4cf30532f98974a8396d885650f3ba5a13026e/numpy-2.5.2-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:c081cbe16ba1ab53078e5ff29013621e33c509eedab055775d956427712c236e", size = 12025331, upload-time = "2026-08-09T13:47:52.646Z" }, + { url = "https://files.pythonhosted.org/packages/8b/3a/2898e003a5fbaf87e76c039b4ee1f5eb390471b4ffe74887c1f34c4e791e/numpy-2.5.2-cp315-cp315t-macosx_14_0_arm64.whl", hash = "sha256:0090ccdd57ec2703e9b49d0bf554767370581c1dd0a6b2bb2b2d9def317d042a", size = 5472336, upload-time = "2026-08-09T13:47:55.403Z" }, + { url = "https://files.pythonhosted.org/packages/61/a5/23f69d07c544597b29758b31b55c27dc9d541012a2c1496189fef702aec2/numpy-2.5.2-cp315-cp315t-macosx_14_0_x86_64.whl", hash = "sha256:6a9bb119fb8dd21ba30b3f0e555b7e2b081bd9883af21ec9c1c633d161cda3a8", size = 6788387, upload-time = "2026-08-09T13:47:58.192Z" }, + { url = "https://files.pythonhosted.org/packages/15/ea/c0dbdbcf22f43782510a3e492dd3da73c6112b69cac8929d16d127536fc4/numpy-2.5.2-cp315-cp315t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a839318485284a6fb31be4f8f2c91c8f2cb22f4543c4a8903f12b0671ffe07cc", size = 15667096, upload-time = "2026-08-09T13:48:01.562Z" }, + { url = "https://files.pythonhosted.org/packages/fc/5e/29c73c31748cdb0f7566642125ba17fd5b56780cddf891b085dab27e4466/numpy-2.5.2-cp315-cp315t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba0a474801b8dc67b66bf465548abc90e82b44d2611b5770f33008dcabffe8ec", size = 16751730, upload-time = "2026-08-09T13:48:05.706Z" }, + { url = "https://files.pythonhosted.org/packages/47/95/02501e8454796bb58dadf7a99d3181e0b464bf264e1003039572f9779fac/numpy-2.5.2-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:0a4035ae1129ff8777f08bfbd44f1e5d8e9c049ce0c2dd78fc0d92c13e7251c0", size = 17038686, upload-time = "2026-08-09T13:48:09.627Z" }, + { url = "https://files.pythonhosted.org/packages/0e/b5/53a681d91b5c82687067d8ea5035e02d917b5509d6f334cb06484a954714/numpy-2.5.2-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:77843ca236b777e67f8d6b3660ea116e499612703a0ecd7093f316201eb9d8e2", size = 18507727, upload-time = "2026-08-09T13:48:13.744Z" }, + { url = "https://files.pythonhosted.org/packages/42/06/6e11443f7b64ee376c860506091103bf68f92d2cab9e8d96d4501babf07c/numpy-2.5.2-cp315-cp315t-win32.whl", hash = "sha256:7354826bc6f8f69402e9b7fe28d15fcd34feebd74f856f111585c5b0c9fb0251", size = 6269775, upload-time = "2026-08-09T13:48:17.543Z" }, + { url = "https://files.pythonhosted.org/packages/f1/18/195d6b86cd72dbbc501edfa778005fa6b87afd34c153e46028cd3a0938f4/numpy-2.5.2-cp315-cp315t-win_amd64.whl", hash = "sha256:e5651f3f87add730ee6608d915009e19c911fba0cb000c7e3ea994b7d768eb12", size = 12782559, upload-time = "2026-08-09T13:48:21.023Z" }, + { url = "https://files.pythonhosted.org/packages/b4/07/458c344f0f0c178f4481dad5cca790626ffe4c34eabf9467069d06ee4999/numpy-2.5.2-cp315-cp315t-win_arm64.whl", hash = "sha256:5f8e00be2ec6f45f4e8a41a527f68d44a7d96fee92a650e4d8b1326f77f61e6e", size = 10748103, upload-time = "2026-08-09T13:48:24.21Z" }, +] + [[package]] name = "packaging" version = "26.2" @@ -92,6 +143,9 @@ wheels = [ name = "terrain-diffusion" version = "0.1.0" source = { editable = "." } +dependencies = [ + { name = "numpy" }, +] [package.dev-dependencies] dev = [ @@ -100,6 +154,7 @@ dev = [ ] [package.metadata] +requires-dist = [{ name = "numpy", specifier = ">=2.5.2" }] [package.metadata.requires-dev] dev = [ From a9b088cda4c596ce767393976c5966c7f3e20867 Mon Sep 17 00:00:00 2001 From: Zain-Mahmoud Date: Wed, 12 Aug 2026 16:03:54 +0300 Subject: [PATCH 07/10] update classes and tests --- src/terrain_diffusion/inference.py | 115 ++++++++++++++++++++++++++--- tests/test_inference.py | 80 +++++++++++++------- 2 files changed, 158 insertions(+), 37 deletions(-) diff --git a/src/terrain_diffusion/inference.py b/src/terrain_diffusion/inference.py index 4820ab2..70c36eb 100644 --- a/src/terrain_diffusion/inference.py +++ b/src/terrain_diffusion/inference.py @@ -17,26 +17,123 @@ - It runs on the GPU compute node. """ +from __future__ import annotations + from abc import ABC, abstractmethod +from dataclasses import dataclass import numpy as np +PATCH_SIZE = (512, 512) +LATENT_MAP_SIZE = (3, 50, 100) # placeholder + -class TerrainModel(ABC): +@dataclass +class ModelOutput(ABC): @abstractmethod - def generate(self, patch: np.ndarray) -> np.ndarray: - """ - Generate an output patch (C, H, W) from an input patch (C, H, W) using the model - """ + def __init__(self): + raise NotImplementedError + + +@dataclass +class ModelInput(ABC): + @abstractmethod + def __init__(self): + raise NotImplementedError + + +class TerrainModel[InputT: ModelInput, OutputT: ModelOutput](ABC): + @abstractmethod + def predict(self, patch: InputT) -> OutputT: raise NotImplementedError - @classmethod @abstractmethod - def load_model(cls, model_path: str) -> TerrainModel: + def load_weights(self, model_path: str): """ Load a model stored in model_path """ raise NotImplementedError - def __call__(self, patch: np.ndarray) -> np.ndarray: - return self.generate(patch) + def __call__(self, patch: InputT) -> OutputT: + return self.predict(patch) + + +@dataclass +class MockCoreModelInput(ModelInput): + patch: np.ndarray + + def __eq__(self, other: MockCoreModelInput): + return np.array_equal(self.patch, other.patch) + + +@dataclass +class MockCoreModelOutput(ModelOutput): + low_res_grid: np.ndarray + latent_map: np.ndarray + + def __eq__(self, other: MockCoreModelInput): + return np.array_equal(self.low_res_grid, other.low_res_grid) and np.array_equal( + self.latent_map, other.latent_map + ) + + +@dataclass +class MockDecoderModelInput(ModelInput): + latent_map: np.ndarray + + def __eq__(self, other: MockDecoderModelInput): + return np.array_equal(self.latent_map, other.latent_map) + + +@dataclass +class MockDecoderModelOutput(ModelOutput): + full_res_grid: np.ndarray + + def __eq__(self, other: MockCoreModelOutput): + return np.array_equal(self.full_res_grid, other.full_res_grid) + + +class MockCoreModel(TerrainModel[MockCoreModelInput, MockCoreModelOutput]): + weights: np.ndarray + + def predict(self, input: MockCoreModelInput) -> MockCoreModelOutput: + + double = input.patch * 2 + low_res_grid = np.resize(double, (PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8)) + latent_map = np.resize(double, LATENT_MAP_SIZE) + output = MockCoreModelOutput(low_res_grid, latent_map) + + return output + + def load_weights(self, model_path: str): + self.weights = np.ones((3, 4, 5)) + + def __call__(self, input: MockCoreModelInput) -> MockCoreModelOutput: + return self.predict(input) + + +class MockDecoderModel(TerrainModel[MockDecoderModelInput, MockDecoderModelOutput]): + weights: np.ndarray + + def predict(self, input: MockDecoderModelInput) -> MockDecoderModelOutput: + + double = input.latent_map * 2 + full_res_grid = np.resize(double, PATCH_SIZE) + output = MockDecoderModelOutput(full_res_grid) + + return output + + def load_weights(self, model_path: str): + self.weights = np.ones((3, 4, 5)) + + def __call__(self, input: MockDecoderModel) -> MockDecoderModel: + return self.predict(input) + + +def load_model(model_name: str) -> TerrainModel: + if model_name == "decoder": + return MockDecoderModel() + elif model_name == "core": + return MockCoreModel() + else: + raise ValueError("Invalid model") diff --git a/tests/test_inference.py b/tests/test_inference.py index 5ef4fde..b32f210 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -6,42 +6,66 @@ import numpy as np import pytest -from terrain_diffusion.inference import TerrainModel +from terrain_diffusion.inference import ( + LATENT_MAP_SIZE, + PATCH_SIZE, + MockCoreModel, + MockCoreModelInput, + MockCoreModelOutput, + MockDecoderModel, + MockDecoderModelInput, + MockDecoderModelOutput, + load_model, +) -class MockTerrainModel(TerrainModel): - """ - Mock model implementing the TerrainModel abstract class - """ +class TestTerrainModel: + @pytest.fixture + def initial_decoder(self) -> MockDecoderModel: + return MockDecoderModel() - def generate(self, patch: np.ndarray) -> np.ndarray: - """ - Generate an output patch (C, H, W) from the input patch - by multiplying by 2 - """ - return patch * 2 + @pytest.fixture + def initial_core(self) -> MockCoreModel: + return MockCoreModel() - @classmethod - def load_model(cls, model_path): - print(f"loaded model: {model_path}") - return MockTerrainModel() + def test_load_model(self): + assert isinstance(load_model("decoder"), MockDecoderModel), ( + "model does not load correct decoder" + ) + assert isinstance(load_model("core"), MockCoreModel), "model does not load correct core" - def __call__(self, patch: np.ndarray) -> np.ndarray: - return self.generate(patch) + def test_predict_core(self, initial_core): + input = MockCoreModelInput(np.ones(PATCH_SIZE)) + actual = initial_core(input) + actual_2 = initial_core(input) -class TestTerrainModel: - @pytest.fixture - def initial_model(self) -> MockTerrainModel: - return MockTerrainModel() + double = 2 * input.patch + + expected_low_res = np.resize(double, (PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8)) + expected_latent = np.resize(double, LATENT_MAP_SIZE) + expected = MockCoreModelOutput(expected_low_res, expected_latent) + + assert actual == expected, "predictions are not equal" + assert actual == actual_2, "model return different predictions on same input" + assert actual.low_res_grid.shape == (PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8), ( + "low resolution map shape is not patch size // 8" + ) + assert actual.latent_map.shape == LATENT_MAP_SIZE, "latent map size is not correct" + + def test_predict_decoder(self, initial_decoder): + input = MockDecoderModelInput(np.ones(LATENT_MAP_SIZE)) - def test_load(self): - assert isinstance(MockTerrainModel.load_model("my_model"), TerrainModel) + actual = initial_decoder(input) + actual_2 = initial_decoder(input) - def test_generate(self, initial_model): - input_patch = np.ones((3, 5, 10)) # 10x5 image with 3 channels + double = 2 * input.latent_map - actual = initial_model(input_patch) - expected = 2 * input_patch + expected_full_res = np.resize(double, PATCH_SIZE) + expected = MockDecoderModelOutput(expected_full_res) - assert np.array_equal(actual, expected), "patches are not equal" + assert actual == expected, "predictions are not equal" + assert actual == actual_2, "model return different predictions on same input" + assert actual.full_res_grid.shape == PATCH_SIZE, ( + "full resolution map does not match patch size" + ) From 1ea31b235f48da0f4fc87440e755c1bcb6ce39c6 Mon Sep 17 00:00:00 2001 From: Zain-Mahmoud Date: Wed, 12 Aug 2026 23:55:25 +0300 Subject: [PATCH 08/10] update logic --- src/terrain_diffusion/inference.py | 34 ++++++++++++++++++------------ tests/test_inference.py | 20 ++++++++++-------- 2 files changed, 31 insertions(+), 23 deletions(-) diff --git a/src/terrain_diffusion/inference.py b/src/terrain_diffusion/inference.py index 70c36eb..25e6dc5 100644 --- a/src/terrain_diffusion/inference.py +++ b/src/terrain_diffusion/inference.py @@ -54,14 +54,14 @@ def load_weights(self, model_path: str): """ raise NotImplementedError - def __call__(self, patch: InputT) -> OutputT: - return self.predict(patch) - @dataclass class MockCoreModelInput(ModelInput): patch: np.ndarray + def __post_init__(self): + assert self.patch.shape == PATCH_SIZE, "invalid input patch shape" + def __eq__(self, other: MockCoreModelInput): return np.array_equal(self.patch, other.patch) @@ -71,6 +71,12 @@ class MockCoreModelOutput(ModelOutput): low_res_grid: np.ndarray latent_map: np.ndarray + def __post_init__(self): + assert self.latent_map.shape == LATENT_MAP_SIZE, "invalid latent map size" + assert self.low_res_grid.shape == (PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8), ( + "invalid low resolution grid shape" + ) + def __eq__(self, other: MockCoreModelInput): return np.array_equal(self.low_res_grid, other.low_res_grid) and np.array_equal( self.latent_map, other.latent_map @@ -81,6 +87,9 @@ def __eq__(self, other: MockCoreModelInput): class MockDecoderModelInput(ModelInput): latent_map: np.ndarray + def __post_init__(self): + assert self.latent_map.shape == LATENT_MAP_SIZE, "invalid latent map size" + def __eq__(self, other: MockDecoderModelInput): return np.array_equal(self.latent_map, other.latent_map) @@ -89,6 +98,9 @@ def __eq__(self, other: MockDecoderModelInput): class MockDecoderModelOutput(ModelOutput): full_res_grid: np.ndarray + def __post_init__(self): + assert self.full_res_grid.shape == PATCH_SIZE, "invalid full resolution grid size" + def __eq__(self, other: MockCoreModelOutput): return np.array_equal(self.full_res_grid, other.full_res_grid) @@ -108,9 +120,6 @@ def predict(self, input: MockCoreModelInput) -> MockCoreModelOutput: def load_weights(self, model_path: str): self.weights = np.ones((3, 4, 5)) - def __call__(self, input: MockCoreModelInput) -> MockCoreModelOutput: - return self.predict(input) - class MockDecoderModel(TerrainModel[MockDecoderModelInput, MockDecoderModelOutput]): weights: np.ndarray @@ -126,14 +135,11 @@ def predict(self, input: MockDecoderModelInput) -> MockDecoderModelOutput: def load_weights(self, model_path: str): self.weights = np.ones((3, 4, 5)) - def __call__(self, input: MockDecoderModel) -> MockDecoderModel: - return self.predict(input) + +MODELS = {"decoder": MockDecoderModel, "core": MockCoreModel} def load_model(model_name: str) -> TerrainModel: - if model_name == "decoder": - return MockDecoderModel() - elif model_name == "core": - return MockCoreModel() - else: - raise ValueError("Invalid model") + if model_name not in MODELS: + raise ValueError("invalid model name") + return MODELS[model_name]() diff --git a/tests/test_inference.py b/tests/test_inference.py index b32f210..468086c 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -37,13 +37,15 @@ def test_load_model(self): def test_predict_core(self, initial_core): input = MockCoreModelInput(np.ones(PATCH_SIZE)) - actual = initial_core(input) - actual_2 = initial_core(input) + actual = initial_core.predict(input) + actual_2 = initial_core.predict(input) - double = 2 * input.patch + expected_low_res = np.ndarray((PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8)) + expected_latent = np.ndarray(LATENT_MAP_SIZE) + + expected_low_res.fill(2) + expected_latent.fill(2) - expected_low_res = np.resize(double, (PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8)) - expected_latent = np.resize(double, LATENT_MAP_SIZE) expected = MockCoreModelOutput(expected_low_res, expected_latent) assert actual == expected, "predictions are not equal" @@ -56,12 +58,12 @@ def test_predict_core(self, initial_core): def test_predict_decoder(self, initial_decoder): input = MockDecoderModelInput(np.ones(LATENT_MAP_SIZE)) - actual = initial_decoder(input) - actual_2 = initial_decoder(input) + actual = initial_decoder.predict(input) + actual_2 = initial_decoder.predict(input) - double = 2 * input.latent_map + expected_full_res = np.ndarray(PATCH_SIZE) + expected_full_res.fill(2) - expected_full_res = np.resize(double, PATCH_SIZE) expected = MockDecoderModelOutput(expected_full_res) assert actual == expected, "predictions are not equal" From f5cda68bc77988cee7853c46db4dfad085f56d3e Mon Sep 17 00:00:00 2001 From: Zain-Mahmoud Date: Thu, 13 Aug 2026 12:13:16 +0300 Subject: [PATCH 09/10] update class structure --- src/terrain_diffusion/inference.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/terrain_diffusion/inference.py b/src/terrain_diffusion/inference.py index 25e6dc5..4fa7429 100644 --- a/src/terrain_diffusion/inference.py +++ b/src/terrain_diffusion/inference.py @@ -21,6 +21,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass +from typing import ClassVar import numpy as np @@ -58,9 +59,10 @@ def load_weights(self, model_path: str): @dataclass class MockCoreModelInput(ModelInput): patch: np.ndarray + patch_shape: ClassVar[tuple] = PATCH_SIZE def __post_init__(self): - assert self.patch.shape == PATCH_SIZE, "invalid input patch shape" + assert self.patch.shape == self.patch_shape, "invalid input patch shape" def __eq__(self, other: MockCoreModelInput): return np.array_equal(self.patch, other.patch) @@ -70,10 +72,12 @@ def __eq__(self, other: MockCoreModelInput): class MockCoreModelOutput(ModelOutput): low_res_grid: np.ndarray latent_map: np.ndarray + low_res_grid_shape: ClassVar[tuple] = (PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8) + latent_map_shape: ClassVar[tuple] = LATENT_MAP_SIZE def __post_init__(self): - assert self.latent_map.shape == LATENT_MAP_SIZE, "invalid latent map size" - assert self.low_res_grid.shape == (PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8), ( + assert self.latent_map.shape == self.latent_map_shape, "invalid latent map size" + assert self.low_res_grid.shape == self.low_res_grid_shape, ( "invalid low resolution grid shape" ) @@ -86,9 +90,10 @@ def __eq__(self, other: MockCoreModelInput): @dataclass class MockDecoderModelInput(ModelInput): latent_map: np.ndarray + latent_map_shape: ClassVar[tuple] = LATENT_MAP_SIZE def __post_init__(self): - assert self.latent_map.shape == LATENT_MAP_SIZE, "invalid latent map size" + assert self.latent_map.shape == self.latent_map_shape, "invalid latent map size" def __eq__(self, other: MockDecoderModelInput): return np.array_equal(self.latent_map, other.latent_map) @@ -97,9 +102,12 @@ def __eq__(self, other: MockDecoderModelInput): @dataclass class MockDecoderModelOutput(ModelOutput): full_res_grid: np.ndarray + full_res_grid_shape: ClassVar[tuple] = PATCH_SIZE def __post_init__(self): - assert self.full_res_grid.shape == PATCH_SIZE, "invalid full resolution grid size" + assert self.full_res_grid.shape == self.full_res_grid_shape, ( + "invalid full resolution grid size" + ) def __eq__(self, other: MockCoreModelOutput): return np.array_equal(self.full_res_grid, other.full_res_grid) From 0e958957763b5d75f202d2f50cde18e233b6e477 Mon Sep 17 00:00:00 2001 From: Zain-Mahmoud Date: Thu, 13 Aug 2026 12:13:30 +0300 Subject: [PATCH 10/10] add fail tests --- tests/test_inference.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_inference.py b/tests/test_inference.py index 468086c..9b3b8c3 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -28,6 +28,24 @@ def initial_decoder(self) -> MockDecoderModel: def initial_core(self) -> MockCoreModel: return MockCoreModel() + def test_core_input_generation(self): + with pytest.raises(AssertionError): + MockCoreModelInput(np.ones((1, 1))) + + def test_core_output_generation(self): + with pytest.raises(AssertionError): + MockCoreModelOutput(np.ones((1, 1)), np.ones(LATENT_MAP_SIZE)) + with pytest.raises(AssertionError): + MockCoreModelOutput(np.ones((PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8)), np.ones((1, 1))) + + def test_decoder_input_generation(self): + with pytest.raises(AssertionError): + MockDecoderModelInput(np.ones((1, 1))) + + def test_decoder_output_generation(self): + with pytest.raises(AssertionError): + MockDecoderModelOutput(np.ones((1, 1))) + def test_load_model(self): assert isinstance(load_model("decoder"), MockDecoderModel), ( "model does not load correct decoder"