Compare commits

..

No commits in common. "638440feff3fd548a1605011e547793ada54b86a" and "b681f0b2c662b62d8dfee3541fa2ff27464fe4e2" have entirely different histories.

11 changed files with 94 additions and 376 deletions

View File

@ -15,7 +15,6 @@ dependencies = [
"h5py>=3.12.1", "h5py>=3.12.1",
"pyyaml>=6.0.2", "pyyaml>=6.0.2",
"pydantic>=2.10.6", "pydantic>=2.10.6",
"torch"
] ]
[project.scripts] [project.scripts]

223
uv.lock
View File

@ -1,9 +1,5 @@
version = 1 version = 1
requires-python = ">=3.11" requires-python = ">=3.11"
resolution-markers = [
"python_full_version >= '3.12'",
"python_full_version < '3.12'",
]
[[package]] [[package]]
name = "annotated-types" name = "annotated-types"
@ -104,15 +100,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321 }, { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321 },
] ]
[[package]]
name = "filelock"
version = "3.17.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/dc/9c/0b15fb47b464e1b663b1acd1253a062aa5feecb07d4e597daea542ebd2b5/filelock-3.17.0.tar.gz", hash = "sha256:ee4e77401ef576ebb38cd7f13b9b28893194acc20a8e68e18730ba9c0e54660e", size = 18027 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/89/ec/00d68c4ddfedfe64159999e5f8a98fb8442729a63e2077eb9dcd89623d27/filelock-3.17.0-py3-none-any.whl", hash = "sha256:533dc2f7ba78dc2f0f531fc6c4940addf7b70a481e269a5a3b93be94ffbe8338", size = 16164 },
]
[[package]] [[package]]
name = "flask" name = "flask"
version = "3.1.0" version = "3.1.0"
@ -175,15 +162,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/1e/6a/6afc55d75036b8d3fe5ceaea2e8da2c04e8f3b298325de73a35f098cb9a8/fonttools-4.55.6-py3-none-any.whl", hash = "sha256:d20ab5a78d0536c26628eaadba661e7ae2427b1e5c748a0a510a44d914e1b155", size = 1112524 }, { url = "https://files.pythonhosted.org/packages/1e/6a/6afc55d75036b8d3fe5ceaea2e8da2c04e8f3b298325de73a35f098cb9a8/fonttools-4.55.6-py3-none-any.whl", hash = "sha256:d20ab5a78d0536c26628eaadba661e7ae2427b1e5c748a0a510a44d914e1b155", size = 1112524 },
] ]
[[package]]
name = "fsspec"
version = "2024.12.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/ee/11/de70dee31455c546fbc88301971ec03c328f3d1138cfba14263f651e9551/fsspec-2024.12.0.tar.gz", hash = "sha256:670700c977ed2fb51e0d9f9253177ed20cbde4a3e5c0283cc5385b5870c8533f", size = 291600 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/de/86/5486b0188d08aa643e127774a99bac51ffa6cf343e3deb0583956dca5b22/fsspec-2024.12.0-py3-none-any.whl", hash = "sha256:b520aed47ad9804237ff878b504267a3b0b441e97508bd6d2d8774e3db85cee2", size = 183862 },
]
[[package]] [[package]]
name = "h11" name = "h11"
version = "0.14.0" version = "0.14.0"
@ -427,24 +405,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979 }, { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979 },
] ]
[[package]]
name = "mpmath"
version = "1.3.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198 },
]
[[package]]
name = "networkx"
version = "3.4.2"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/fd/1d/06475e1cd5264c0b870ea2cc6fdb3e37177c1e565c43f56ff17a10e3937f/networkx-3.4.2.tar.gz", hash = "sha256:307c3669428c5362aab27c8a1260aa8f47c4e91d3891f48be0141738d8d053e1", size = 2151368 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/b9/54/dd730b32ea14ea797530a4479b2ed46a6fb250f682a9cfb997e968bf0261/networkx-3.4.2-py3-none-any.whl", hash = "sha256:df5d4365b724cf81b8c6a7312509d0c22386097011ad1abe274afd5e9d3bbc5f", size = 1723263 },
]
[[package]] [[package]]
name = "numpy" name = "numpy"
version = "2.2.2" version = "2.2.2"
@ -493,116 +453,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/80/94/cd9e9b04012c015cb6320ab3bf43bc615e248dddfeb163728e800a5d96f0/numpy-2.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:97b974d3ba0fb4612b77ed35d7627490e8e3dff56ab41454d9e8b23448940576", size = 12696208 }, { url = "https://files.pythonhosted.org/packages/80/94/cd9e9b04012c015cb6320ab3bf43bc615e248dddfeb163728e800a5d96f0/numpy-2.2.2-cp313-cp313t-win_amd64.whl", hash = "sha256:97b974d3ba0fb4612b77ed35d7627490e8e3dff56ab41454d9e8b23448940576", size = 12696208 },
] ]
[[package]]
name = "nvidia-cublas-cu12"
version = "12.4.5.8"
source = { registry = "https://pypi.org/simple" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/ae/71/1c91302526c45ab494c23f61c7a84aa568b8c1f9d196efa5993957faf906/nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl", hash = "sha256:2fc8da60df463fdefa81e323eef2e36489e1c94335b5358bcb38360adf75ac9b", size = 363438805 },
]
[[package]]
name = "nvidia-cuda-cupti-cu12"
version = "12.4.127"
source = { registry = "https://pypi.org/simple" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/67/42/f4f60238e8194a3106d06a058d494b18e006c10bb2b915655bd9f6ea4cb1/nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:9dec60f5ac126f7bb551c055072b69d85392b13311fcc1bcda2202d172df30fb", size = 13813957 },
]
[[package]]
name = "nvidia-cuda-nvrtc-cu12"
version = "12.4.127"
source = { registry = "https://pypi.org/simple" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/2c/14/91ae57cd4db3f9ef7aa99f4019cfa8d54cb4caa7e00975df6467e9725a9f/nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a178759ebb095827bd30ef56598ec182b85547f1508941a3d560eb7ea1fbf338", size = 24640306 },
]
[[package]]
name = "nvidia-cuda-runtime-cu12"
version = "12.4.127"
source = { registry = "https://pypi.org/simple" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/ea/27/1795d86fe88ef397885f2e580ac37628ed058a92ed2c39dc8eac3adf0619/nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:64403288fa2136ee8e467cdc9c9427e0434110899d07c779f25b5c068934faa5", size = 883737 },
]
[[package]]
name = "nvidia-cudnn-cu12"
version = "9.1.0.70"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "nvidia-cublas-cu12" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/9f/fd/713452cd72343f682b1c7b9321e23829f00b842ceaedcda96e742ea0b0b3/nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f", size = 664752741 },
]
[[package]]
name = "nvidia-cufft-cu12"
version = "11.2.1.3"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "nvidia-nvjitlink-cu12" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/27/94/3266821f65b92b3138631e9c8e7fe1fb513804ac934485a8d05776e1dd43/nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f083fc24912aa410be21fa16d157fed2055dab1cc4b6934a0e03cba69eb242b9", size = 211459117 },
]
[[package]]
name = "nvidia-curand-cu12"
version = "10.3.5.147"
source = { registry = "https://pypi.org/simple" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/8a/6d/44ad094874c6f1b9c654f8ed939590bdc408349f137f9b98a3a23ccec411/nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a88f583d4e0bb643c49743469964103aa59f7f708d862c3ddb0fc07f851e3b8b", size = 56305206 },
]
[[package]]
name = "nvidia-cusolver-cu12"
version = "11.6.1.9"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "nvidia-cublas-cu12" },
{ name = "nvidia-cusparse-cu12" },
{ name = "nvidia-nvjitlink-cu12" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/3a/e1/5b9089a4b2a4790dfdea8b3a006052cfecff58139d5a4e34cb1a51df8d6f/nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl", hash = "sha256:19e33fa442bcfd085b3086c4ebf7e8debc07cfe01e11513cc6d332fd918ac260", size = 127936057 },
]
[[package]]
name = "nvidia-cusparse-cu12"
version = "12.3.1.170"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "nvidia-nvjitlink-cu12" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/db/f7/97a9ea26ed4bbbfc2d470994b8b4f338ef663be97b8f677519ac195e113d/nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl", hash = "sha256:ea4f11a2904e2a8dc4b1833cc1b5181cde564edd0d5cd33e3c168eff2d1863f1", size = 207454763 },
]
[[package]]
name = "nvidia-nccl-cu12"
version = "2.21.5"
source = { registry = "https://pypi.org/simple" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/df/99/12cd266d6233f47d00daf3a72739872bdc10267d0383508b0b9c84a18bb6/nvidia_nccl_cu12-2.21.5-py3-none-manylinux2014_x86_64.whl", hash = "sha256:8579076d30a8c24988834445f8d633c697d42397e92ffc3f63fa26766d25e0a0", size = 188654414 },
]
[[package]]
name = "nvidia-nvjitlink-cu12"
version = "12.4.127"
source = { registry = "https://pypi.org/simple" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/ff/ff/847841bacfbefc97a00036e0fce5a0f086b640756dc38caea5e1bb002655/nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:06b3b9b25bf3f8af351d664978ca26a16d2c5127dbd53c0497e28d1fb9611d57", size = 21066810 },
]
[[package]]
name = "nvidia-nvtx-cu12"
version = "12.4.127"
source = { registry = "https://pypi.org/simple" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/87/20/199b8713428322a2f22b722c62b8cc278cc53dffa9705d744484b5035ee9/nvidia_nvtx_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:781e950d9b9f60d8241ccea575b32f5105a5baf4c2351cab5256a24869f12a1a", size = 99144 },
]
[[package]] [[package]]
name = "optype" name = "optype"
version = "0.9.0" version = "0.9.0"
@ -896,15 +746,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/71/1f/f1db21491a40923251244ee8de4f71ce95dee2ed83c2b49ad479c2ea5189/scipy_stubs-1.15.1.0-py3-none-any.whl", hash = "sha256:f6e8f8dfd2aaa343705c46ac1fdc69556ea8a0a9b0645cb3af4dd652665a452f", size = 454057 }, { url = "https://files.pythonhosted.org/packages/71/1f/f1db21491a40923251244ee8de4f71ce95dee2ed83c2b49ad479c2ea5189/scipy_stubs-1.15.1.0-py3-none-any.whl", hash = "sha256:f6e8f8dfd2aaa343705c46ac1fdc69556ea8a0a9b0645cb3af4dd652665a452f", size = 454057 },
] ]
[[package]]
name = "setuptools"
version = "75.8.0"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/92/ec/089608b791d210aec4e7f97488e67ab0d33add3efccb83a056cbafe3a2a6/setuptools-75.8.0.tar.gz", hash = "sha256:c5afc8f407c626b8313a86e10311dd3f661c6cd9c09d4bf8c15c0e11f9f2b0e6", size = 1343222 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/69/8a/b9dc7678803429e4a3bc9ba462fa3dd9066824d3c607490235c6a796be5a/setuptools-75.8.0-py3-none-any.whl", hash = "sha256:e3982f444617239225d675215d51f6ba05f845d4eec313da4418fdbb56fb27e3", size = 1228782 },
]
[[package]] [[package]]
name = "shellingham" name = "shellingham"
version = "1.5.4" version = "1.5.4"
@ -935,68 +776,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050 }, { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050 },
] ]
[[package]]
name = "sympy"
version = "1.13.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "mpmath" },
]
sdist = { url = "https://files.pythonhosted.org/packages/ca/99/5a5b6f19ff9f083671ddf7b9632028436167cd3d33e11015754e41b249a4/sympy-1.13.1.tar.gz", hash = "sha256:9cebf7e04ff162015ce31c9c6c9144daa34a93bd082f54fd8f12deca4f47515f", size = 7533040 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/b2/fe/81695a1aa331a842b582453b605175f419fe8540355886031328089d840a/sympy-1.13.1-py3-none-any.whl", hash = "sha256:db36cdc64bf61b9b24578b6f7bab1ecdd2452cf008f34faa33776680c26d66f8", size = 6189177 },
]
[[package]]
name = "torch"
version = "2.5.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "filelock" },
{ name = "fsspec" },
{ name = "jinja2" },
{ name = "networkx" },
{ name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "setuptools", marker = "python_full_version >= '3.12'" },
{ name = "sympy" },
{ name = "triton", marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux'" },
{ name = "typing-extensions" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/d1/35/e8b2daf02ce933e4518e6f5682c72fd0ed66c15910ea1fb4168f442b71c4/torch-2.5.1-cp311-cp311-manylinux1_x86_64.whl", hash = "sha256:de5b7d6740c4b636ef4db92be922f0edc425b65ed78c5076c43c42d362a45457", size = 906474467 },
{ url = "https://files.pythonhosted.org/packages/40/04/bd91593a4ca178ece93ca55f27e2783aa524aaccbfda66831d59a054c31e/torch-2.5.1-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:340ce0432cad0d37f5a31be666896e16788f1adf8ad7be481196b503dad675b9", size = 91919450 },
{ url = "https://files.pythonhosted.org/packages/0d/4a/e51420d46cfc90562e85af2fee912237c662ab31140ab179e49bd69401d6/torch-2.5.1-cp311-cp311-win_amd64.whl", hash = "sha256:603c52d2fe06433c18b747d25f5c333f9c1d58615620578c326d66f258686f9a", size = 203098237 },
{ url = "https://files.pythonhosted.org/packages/d0/db/5d9cbfbc7968d79c5c09a0bc0bc3735da079f2fd07cc10498a62b320a480/torch-2.5.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:31f8c39660962f9ae4eeec995e3049b5492eb7360dd4f07377658ef4d728fa4c", size = 63884466 },
{ url = "https://files.pythonhosted.org/packages/8b/5c/36c114d120bfe10f9323ed35061bc5878cc74f3f594003854b0ea298942f/torch-2.5.1-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:ed231a4b3a5952177fafb661213d690a72caaad97d5824dd4fc17ab9e15cec03", size = 906389343 },
{ url = "https://files.pythonhosted.org/packages/6d/69/d8ada8b6e0a4257556d5b4ddeb4345ea8eeaaef3c98b60d1cca197c7ad8e/torch-2.5.1-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:3f4b7f10a247e0dcd7ea97dc2d3bfbfc90302ed36d7f3952b0008d0df264e697", size = 91811673 },
{ url = "https://files.pythonhosted.org/packages/5f/ba/607d013b55b9fd805db2a5c2662ec7551f1910b4eef39653eeaba182c5b2/torch-2.5.1-cp312-cp312-win_amd64.whl", hash = "sha256:73e58e78f7d220917c5dbfad1a40e09df9929d3b95d25e57d9f8558f84c9a11c", size = 203046841 },
{ url = "https://files.pythonhosted.org/packages/57/6c/bf52ff061da33deb9f94f4121fde7ff3058812cb7d2036c97bc167793bd1/torch-2.5.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:8c712df61101964eb11910a846514011f0b6f5920c55dbf567bff8a34163d5b1", size = 63858109 },
{ url = "https://files.pythonhosted.org/packages/69/72/20cb30f3b39a9face296491a86adb6ff8f1a47a897e4d14667e6cf89d5c3/torch-2.5.1-cp313-cp313-manylinux1_x86_64.whl", hash = "sha256:9b61edf3b4f6e3b0e0adda8b3960266b9009d02b37555971f4d1c8f7a05afed7", size = 906393265 },
]
[[package]]
name = "triton"
version = "3.1.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "filelock" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/86/17/d9a5cf4fcf46291856d1e90762e36cbabd2a56c7265da0d1d9508c8e3943/triton-3.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0f34f6e7885d1bf0eaaf7ba875a5f0ce6f3c13ba98f9503651c1e6dc6757ed5c", size = 209506424 },
{ url = "https://files.pythonhosted.org/packages/78/eb/65f5ba83c2a123f6498a3097746607e5b2f16add29e36765305e4ac7fdd8/triton-3.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c8182f42fd8080a7d39d666814fa36c5e30cc00ea7eeeb1a2983dbb4c99a0fdc", size = 209551444 },
]
[[package]] [[package]]
name = "typer" name = "typer"
version = "0.15.1" version = "0.15.1"
@ -1048,7 +827,6 @@ dependencies = [
{ name = "pyyaml" }, { name = "pyyaml" },
{ name = "scipy" }, { name = "scipy" },
{ name = "scipy-stubs" }, { name = "scipy-stubs" },
{ name = "torch" },
{ name = "typer" }, { name = "typer" },
] ]
@ -1064,7 +842,6 @@ requires-dist = [
{ name = "pyyaml", specifier = ">=6.0.2" }, { name = "pyyaml", specifier = ">=6.0.2" },
{ name = "scipy" }, { name = "scipy" },
{ name = "scipy-stubs" }, { name = "scipy-stubs" },
{ name = "torch" },
{ name = "typer" }, { name = "typer" },
] ]

View File

@ -3,7 +3,6 @@ import multiprocessing as mp
import numpy as np import numpy as np
import numpy.typing as npt import numpy.typing as npt
import torch
import typer import typer
from .. import visualise from .. import visualise
@ -14,7 +13,6 @@ from . import file, globals
app = typer.Typer(callback=globals.main) app = typer.Typer(callback=globals.main)
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@app.command() @app.command()
@ -38,25 +36,19 @@ def antennas() -> None:
def heatmap() -> None: def heatmap() -> None:
preprocessor = Preprocessor() preprocessor = Preprocessor()
aoa = AoA() aoa = AoA()
# Start webapp in background process
manager = mp.Manager() manager = mp.Manager()
webapp_queue: "mp.Queue[npt.NDArray[np.float32]]" = manager.Queue( webapp_queue: "mp.Queue[AoA]" = manager.Queue(config.sample_rate)
config.sample_rate # Start webapp in background process
)
webapp = mp.Process(target=visualise.start, args=(webapp_queue,)) webapp = mp.Process(target=visualise.start, args=(webapp_queue,))
webapp.start() webapp.start()
def callback(antenna_data: npt.NDArray[np.complex64]) -> None: def callback(antenna_data: npt.NDArray[np.complex128]) -> None:
logger.info(f"Got final CSI data with shape {antenna_data.shape}") logger.info(f"Got final CSI data with shape {antenna_data.shape}")
processed = preprocessor.preprocess(antenna_data) processed = preprocessor.preprocess(antenna_data)
logger.info(f"Processed CSI data with shape {processed.shape}") # visualise.add_data(all_data, processed)
processed_tensor = torch.tensor(processed, device=device) aoa.update(processed)
aoa.update(processed_tensor)
if not webapp_queue.full(): if not webapp_queue.full():
heatmap = aoa.heatmap() webapp_queue.put(aoa)
webapp_queue.put(heatmap)
globals.csi_producer(csi_callback=callback) globals.csi_producer(csi_callback=callback)
logger.info("Finished processing CSI data") logger.info("Finished processing CSI data")

View File

@ -93,14 +93,14 @@ class CSIHeader:
class CSI: class CSI:
@staticmethod @staticmethod
def parseCsiData(data: bytes, header: CSIHeader) -> npt.NDArray[np.complex64]: def parseCsiData(data: bytes, header: CSIHeader) -> npt.NDArray[np.complex128]:
csi_matrix: npt.NDArray[np.complex64] = np.zeros( csi_matrix: npt.NDArray[np.complex128] = np.zeros(
( (
header.num_subcarriers, header.num_subcarriers,
header.num_rx, header.num_rx,
header.num_tx, header.num_tx,
), ),
dtype=np.complex64, dtype=np.complex128,
) )
pos = 0 pos = 0
for j in range(header.num_rx): for j in range(header.num_rx):

View File

@ -3,7 +3,7 @@ from typing import Callable, Protocol
import numpy as np import numpy as np
import numpy.typing as npt import numpy.typing as npt
CSICallback = Callable[[npt.NDArray[np.complex64]], None] CSICallback = Callable[[npt.NDArray[np.complex128]], None]
class CSIProducer(Protocol): class CSIProducer(Protocol):

View File

@ -23,13 +23,6 @@ class MUSIC(BaseModel):
eigval_threshold: int eigval_threshold: int
window_size: int window_size: int
class Heatmap(BaseModel):
theta_resolution: int
tof_resolution: int
tof_max: float
heatmap: Heatmap
class Antennas(BaseModel): class Antennas(BaseModel):
spacing: float spacing: float

View File

@ -3,25 +3,21 @@ from datetime import datetime
import numpy as np import numpy as np
import numpy.typing as npt import numpy.typing as npt
import torch
from ..config import config from ..config import config
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.set_default_device(device)
class AoA: class AoA:
def __init__(self) -> None: def __init__(self) -> None:
self.historical_autocorr = torch.tensor([], dtype=torch.complex64) self.historical_autocorr = np.array([])
self.N_subcarriers = -1 self.N_subcarriers = -1
self.N_rx = -1 self.N_rx = -1
self.timestamp = datetime.now() self.timestamp = datetime.now()
pass pass
def smooth(self, data: torch.Tensor) -> torch.Tensor: def smooth(self, data: npt.NDArray[np.complex128]) -> npt.NDArray[np.complex128]:
assert len(data.shape) == 3 assert len(data.shape) == 3
M = data.shape[0] # Number of subcarriers M = data.shape[0] # Number of subcarriers
@ -36,151 +32,99 @@ class AoA:
# This only works with 1 TX antenna (i.e. no MIMO) - see #4 for more details # This only works with 1 TX antenna (i.e. no MIMO) - see #4 for more details
assert T == 1, "The current implementation only supports 1 TX antenna" assert T == 1, "The current implementation only supports 1 TX antenna"
H_n = torch.zeros((N, M // 2, M // 2 + 1), dtype=torch.complex64) H_n = np.zeros((N, M // 2, M // 2 + 1), dtype=np.complex128)
for i in range(N): for i in range(N):
for j in range(M // 2): for j in range(M // 2):
H_n[i, j] = data[j : j + M // 2 + 1, i, 0] H_n[i, j] = data[j : j + M // 2 + 1, i, 0]
H_sm_rows = [torch.hstack(list(H_n[i : i + N // 2 + 1])) for i in range(N // 2)] H_sm_rows = [np.hstack(H_n[i : i + N // 2 + 1]) for i in range(N // 2)]
H_sm = torch.vstack(H_sm_rows) H_sm = np.vstack(H_sm_rows)
logger.debug(f"Smoothed: {H_sm.shape}") logger.debug(f"Smoothed: {H_sm.shape}")
return H_sm return H_sm
def update(self, data: torch.Tensor) -> None: def update(self, data: npt.NDArray[np.complex128]) -> None:
self.timestamp = datetime.now() self.timestamp = datetime.now()
H_sm = self.smooth(data) H_sm = self.smooth(data)
logger.debug(f"Calculated smoothed CSI matrix: {H_sm.shape}")
auto_corr = H_sm @ torch.conj(H_sm).T auto_corr = np.matmul(H_sm, np.conj(H_sm).T)
# This matrix is by definition Hermitian. # This matrix is by definition Hermitian.
# Therefore, all of its eigenvectors are orthogonal. # Therefore, all of its eigenvectors are orthogonal.
if len(self.historical_autocorr.shape) <= 1: if self.historical_autocorr.size == 0:
self.historical_autocorr = torch.unsqueeze(auto_corr, 0) self.historical_autocorr = np.expand_dims(auto_corr, 0)
else: else:
self.historical_autocorr = torch.cat( self.historical_autocorr = np.append(
(self.historical_autocorr, torch.unsqueeze(auto_corr, 0)) self.historical_autocorr, np.expand_dims(auto_corr, 0), axis=0
) )
WINDOW_SIZE = config.music.window_size WINDOW_SIZE = config.music.window_size
if self.historical_autocorr.shape[0] > WINDOW_SIZE: if self.historical_autocorr.shape[0] > WINDOW_SIZE:
self.historical_autocorr = self.historical_autocorr[-WINDOW_SIZE:] self.historical_autocorr = self.historical_autocorr[-WINDOW_SIZE:]
logger.debug("Finished updating autocorrelation matrix") # Is the moving average also Hermitian?
R = np.mean(self.historical_autocorr, axis=0)
def steering_vector(self, theta: torch.Tensor, tof: torch.Tensor) -> torch.Tensor: # The smallest eigenvectors span the noise subspace,
assert theta.shape == tof.shape # and the largest span the signal subspace.
assert len(theta.shape) == 1 eigvals, eigvecs = np.linalg.eigh(R)
N = theta.shape[0] self.E_n = eigvecs[:, np.abs(eigvals) < config.music.eigval_threshold]
omega_t: torch.Tensor = torch.exp(-2j * np.pi * config.delta_f * tof) def steering_vector(
phi_theta: torch.Tensor = torch.exp( self, theta: float, tof: float
) -> npt.NDArray[np.complexfloating]:
omega_t: npt.NDArray[np.complex128] = np.exp(-2j * np.pi * config.delta_f * tof)
phi_theta: npt.NDArray[np.complex128] = np.exp(
2j 2j
* np.pi * np.pi
* config.central_freq_hz * config.central_freq_hz
* config.antennas.spacing * config.antennas.spacing
* (1 - torch.cos(theta)) * (1 - np.cos(theta))
/ 299_792_458 / 299_792_458
) )
assert omega_t.shape == phi_theta.shape == (N,)
omega_t = torch.unsqueeze(omega_t, dim=-1) omega_t = np.expand_dims(omega_t, axis=-1)
phi_theta = torch.unsqueeze(phi_theta, dim=-1) phi_theta = np.expand_dims(phi_theta, axis=-1)
assert omega_t.shape == phi_theta.shape == (N, 1) antenna_v = omega_t ** np.arange(self.N_subcarriers // 2)
phis = phi_theta ** np.arange(self.N_rx // 2)
antenna_v = np.expand_dims(antenna_v, axis=-1)
steering = antenna_v * phis
return steering.T.reshape(-1)
antenna_v = omega_t ** torch.arange( def evaluate(self, theta: float, tof: float) -> float:
self.N_subcarriers // 2, dtype=torch.float32 try:
) steering = self.steering_vector(theta, tof)
phis = phi_theta ** torch.arange(self.N_rx // 2, dtype=torch.float32) steering_h = np.conj(steering).T
except Exception as e:
assert antenna_v.shape == (N, self.N_subcarriers // 2) logger.exception(e)
assert phis.shape == (N, self.N_rx // 2) return 0
E_n = self.E_n
antenna_v = torch.unsqueeze(antenna_v, dim=1) E_n_H = np.conj(E_n).T
phis = torch.unsqueeze(phis, dim=-1) c = 1 / (0.001 + (steering_h @ E_n @ E_n_H @ steering))
return np.abs(c.real)
assert antenna_v.shape == (N, 1, self.N_subcarriers // 2)
assert phis.shape == (N, self.N_rx // 2, 1)
steering = torch.bmm(phis, antenna_v)
assert steering.shape == (N, self.N_rx // 2, self.N_subcarriers // 2)
return steering.reshape(N, -1)
def evaluate(self, theta: torch.Tensor, tof: torch.Tensor) -> torch.Tensor:
R = torch.mean(self.historical_autocorr, dim=0)
assert (R == torch.conj(R).T).all()
# The smallest eigenvectors span the noise subspace,
# and the largest span the signal subspace.
logger.debug(f"Calculating eigenvectors of R: {R.shape}")
eigvals, eigvecs = torch.linalg.eigh(R)
assert isinstance(eigvals, torch.Tensor)
assert isinstance(eigvecs, torch.Tensor)
logger.info(f"Eigenvalues: {eigvals}")
E_n = eigvecs[:, torch.abs(eigvals) < config.music.eigval_threshold]
logger.debug(f"Signal subspace: {E_n.shape}")
steering = torch.unsqueeze(self.steering_vector(theta, tof), dim=-1)
steering_h = torch.conj(steering).permute(0, 2, 1)
E_n = E_n.unsqueeze(0)
E_n_H = torch.conj(E_n).permute(0, 2, 1)
logger.debug(
f"Heatmap multiplication: {steering_h.shape}, {E_n.shape}, "
f"{E_n_H.shape}, {steering.shape}"
)
c: torch.Tensor = 1 / (steering_h @ E_n @ E_n_H @ steering)
return torch.abs(c)[:, 0, 0]
def heatmap(self) -> npt.NDArray[np.float32]:
thetas = np.linspace(
0, np.pi, config.music.heatmap.theta_resolution, dtype=np.float32
)
tofs = np.linspace(
0,
config.music.heatmap.tof_max,
config.music.heatmap.tof_resolution,
dtype=np.float32,
)
thetas_mesh, tofs_mesh = np.meshgrid(thetas, tofs)
logger.debug(
f"Calculating heatmap with {thetas_mesh.shape} and {tofs_mesh.shape}"
)
evaluated = self.evaluate(
torch.tensor(thetas_mesh.reshape(-1)),
torch.tensor(tofs_mesh.reshape(-1)),
)
logger.debug(f"Evaluated heatmap: {evaluated.shape}")
heatmap: npt.NDArray[np.float32] = evaluated.reshape(
config.music.heatmap.tof_resolution,
config.music.heatmap.theta_resolution,
).numpy(force=True)
return heatmap
def test_smoothing() -> None: def test_smoothing() -> None:
row, col = np.indices((6, 4)) row, col = np.indices((4, 2))
data = row + 1j * col data = row + 1j * col
data = np.expand_dims(data, axis=2)
np.set_printoptions(linewidth=200)
print(data.shape)
aoa = AoA() aoa = AoA()
aoa.N_subcarriers = 6
aoa.N_rx = 4
smoothed = aoa.smooth(data) smoothed = aoa.smooth(data)
print(smoothed) H_0 = np.array([[0 + 0j, 0 + 1j, 0 + 2j], [0 + 1j, 0 + 2j, 0 + 3j]])
H_01 = np.vstack([H_0, H_0 + 1])
H_12 = np.vstack([H_0 + 1, H_0 + 2])
expected = np.hstack([H_01, H_12])
print(expected)
assert np.allclose(smoothed, expected)
def test_steering_vector() -> None: def test_steering_vector() -> None:
aoa = AoA() aoa = AoA()
aoa.N_subcarriers = 10 aoa.N_subcarriers = 10
aoa.N_rx = 4 aoa.N_rx = 2
tau = torch.Tensor([1, 0]) tau = 1
theta = torch.Tensor([0, 1]) theta = 0
print(aoa.steering_vector(theta, tau)) print(aoa.steering_vector(theta, tau))
assert False assert False

View File

@ -15,9 +15,9 @@ np.seterr(invalid="ignore")
class Preprocessor: class Preprocessor:
def __init__(self) -> None: def __init__(self) -> None:
self.prev_entries: Queue[npt.NDArray[np.complex64]] = Queue(maxsize=100) self.prev_entries: Queue[npt.NDArray[np.complex128]] = Queue(maxsize=100)
self.short_term_avg = np.zeros((1,), dtype=np.complex64) self.short_term_avg = np.zeros((1,), dtype=np.complex128)
self.long_term_avg = np.zeros((1,), dtype=np.complex64) self.long_term_avg = np.zeros((1,), dtype=np.complex128)
self.filter = butter( self.filter = butter(
5, 5,
config.preprocessing.bandpass.bounds, config.preprocessing.bandpass.bounds,
@ -26,7 +26,7 @@ class Preprocessor:
output="sos", output="sos",
) )
def preprocess(self, h: npt.NDArray[np.complex64]) -> npt.NDArray[np.complex64]: def preprocess(self, h: npt.NDArray[np.complex128]) -> npt.NDArray[np.complex128]:
# CSI data is not available for pilot subcarriers. # CSI data is not available for pilot subcarriers.
h_hat = np.where( h_hat = np.where(
np.expand_dims(h[:, 0, 0] == 0, axis=(1, 2)), np.expand_dims(h[:, 0, 0] == 0, axis=(1, 2)),
@ -46,7 +46,7 @@ class Preprocessor:
# Assume that all csi matrices will have the same shape # Assume that all csi matrices will have the same shape
if self.long_term_avg.shape != h_hat.shape: if self.long_term_avg.shape != h_hat.shape:
self.long_term_avg = np.zeros(h_hat.shape, dtype=np.complex64) self.long_term_avg = np.zeros(h_hat.shape, dtype=np.complex128)
self.long_term_avg = ( self.long_term_avg = (
self.long_term_avg * (1 - config.preprocessing.moving_average_alpha) self.long_term_avg * (1 - config.preprocessing.moving_average_alpha)

View File

@ -13,7 +13,7 @@ from flask import Flask, Response, render_template, request
from flask_sock import Sock from flask_sock import Sock
from simple_websocket import Server from simple_websocket import Server
from ..config import config from .. import config
from ..processing.aoa import AoA from ..processing.aoa import AoA
matplotlib.use("agg") matplotlib.use("agg")
@ -92,18 +92,18 @@ def add_data(
del subscriber_settings[subscriber] del subscriber_settings[subscriber]
def plot_heatmap(heatmap: npt.NDArray[np.float32]) -> io.BytesIO: def make_heatmap(aoa: AoA, max_tof: float) -> io.BytesIO:
logger.info(f"Making heatmap with aoa of {aoa.timestamp}") logger.info(f"Making heatmap with aoa of {aoa.timestamp}")
fig = plt.figure() fig = plt.figure()
ax = fig.add_axes([0.1, 0.1, 0.9, 0.9]) # , polar=True) ax = fig.add_axes([0, 0, 1, 1], polar=True)
r = np.linspace( r = np.linspace(0, max_tof, 100) # Radius values
0, config.music.heatmap.tof_max, config.music.heatmap.tof_resolution theta = np.linspace(0, np.pi, 50) # Angle values
) R, Theta = np.meshgrid(r, theta) # Create a 2D grid of r and theta
theta = np.linspace(0, np.pi, config.music.heatmap.theta_resolution) # Angle values
mesh = ax.pcolormesh(theta, r, heatmap, edgecolors="face", vmin=0, vmax=50) # Compute the function values
Z = np.log(np.vectorize(aoa.evaluate)(Theta, R))
fig.colorbar(mesh, ax=ax) ax.pcolormesh(Theta, R, Z, edgecolors="face")
buf = io.BytesIO() buf = io.BytesIO()
fig.savefig(buf, format="jpeg") fig.savefig(buf, format="jpeg")
plt.close(fig) plt.close(fig)
@ -112,23 +112,36 @@ def plot_heatmap(heatmap: npt.NDArray[np.float32]) -> io.BytesIO:
return buf return buf
def gather_aoa() -> Generator[bytes, None, None]: def gather_aoa(max_tof: float) -> Generator[bytes, None, None]:
assert aoa_queue is not None assert aoa_queue is not None
prev_frame = datetime.now()
while True: while True:
while (datetime.now() - prev_frame).total_seconds() < 1 / config.HEATMAP_FPS:
time.sleep(0.01)
while not aoa_queue.empty():
logger.debug("Receiving from aoa pipe") logger.debug("Receiving from aoa pipe")
heatmap = aoa_queue.get() aoa = aoa_queue.get()
buf = plot_heatmap(heatmap) prev_frame = datetime.now()
logger.debug(f"Generating heatmap of time {aoa.timestamp}")
buf = make_heatmap(aoa, max_tof)
yield (b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + buf.read() + b"\r\n") yield (b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + buf.read() + b"\r\n")
buf.close() buf.close()
@app.route("/aoa_tof") @app.route("/aoa_tof")
def aoa_tof() -> Response: def aoa_tof() -> Response:
return Response(gather_aoa(), mimetype="multipart/x-mixed-replace; boundary=frame") max_tof_str = request.args.get("max_tof")
try:
max_tof = float(max_tof_str)
except Exception:
max_tof = 5e-8
return Response(
gather_aoa(max_tof), mimetype="multipart/x-mixed-replace; boundary=frame"
)
def start(conn: "mp.Queue[npt.NDArray[np.float32]]") -> None: def start(conn: "mp.Queue[AoA]") -> None:
global app, aoa_queue global app, aoa_queue
aoa_queue = conn aoa_queue = conn