diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..bef6dc343e440fe95b42e61a97873b7ab54fc827 --- /dev/null +++ b/.gitignore @@ -0,0 +1,144 @@ +.vscode +.DS_Store +dask-worker-space +*~ +lightning_logs + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..fef8c29fe0ffed934a19ef1d235e5c1b6ea051fe --- /dev/null +++ b/LICENSE @@ -0,0 +1,202 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..2eb8b71999950f9b5f58ed41f66169b5048f8042 --- /dev/null +++ b/README.md @@ -0,0 +1,106 @@ +LDCast is a precipitation nowcasting model based on a latent diffusion model (LDM, used by e.g. [Stable Diffusion](https://github.com/CompVis/stable-diffusion)). + +This repository contains the code for using LDCast to make predictions and the code used to generate the analysis in the LDCast paper (a preprint is available at https://arxiv.org/abs/2304.12891). + +A GPU is recommended for both using and training LDCast, although you may be able to generate some samples with a CPU and enough patience. + +# Installation + +It is recommended you install the code in its own virtual environment (created with e.g. pyenv or conda). + +Clone the repository, then, in the main directory, run +```bash +$ pip install -e . +``` +This should automatically install the required packages (which might take some minutes). In the paper, we used PyTorch 11.2 but are not aware of any problems with newer versions. + +If you don't want the requirements to be installed (e.g. if you installed them manually with conda), use: +```bash +$ pip install --no-dependencies -e . +``` + +# Using LDCast + +## Pretrained models + +The pretrained models are available at the Zenodo repository https://doi.org/10.5281/zenodo.7780914. Unzip the file `ldcast-models.zip`. The default is to unzip it to the `models` directory, but you can also use another location. + +## Producing predictions + +The easiest way to produce predictions is to use the `ldcast.forecast.Forecast` class, which will set up all models and data transformations and is callable with a past precipitation array. +```python +from ldcast import forecast + +fc = forecast.Forecast( + ldm_weights_fn=ldm_weights_fn, autoenc_weights_fn=autoenc_weights_fn +) +R_pred = fc(R_past) +``` +Here, `ldm_weights_fn` is the path to the LDM weights and `autoenc_weights_fn` is the path to the autoencoder weights. `R_past` is a NumPy array of precipitation rates with shape `(timesteps, height, width)` where `timesteps` must be 4 and `height` and `width` must be divisible by 32. + +### Ensemble predictions + +If want to process multiple cases at once and/or generate several ensemble members, there is the `ldcast.forecast.ForecastDistributed` class. The usage is similar to the `Forecast` class, for example: +```python +from ldcast import forecast + +fc = forecast.ForecastDistributed( + ldm_weights_fn=ldm_weights_fn, autoenc_weights_fn=autoenc_weights_fn +) +R_pred = fc(R_past, ensemble_members=32) +``` +Here, `R_past` should be of shape `(cases, timesteps, height, width)` where `cases` is the number of cases you want to process. For each case, `ensemble_members` predictions are produced (this is the last axis of `R_pred`). `ForecastDistributed` automatically distributes the workload to multiple GPUs if you have them. + +## Demo + +For a practical example, you can run the demo in the `scripts` directory. First download the `ldcast-demo-20210622.zip` file from the [Zenodo repository](https://doi.org/10.5281/zenodo.7780914), then unzip it in the `data` directory. Then run +```bash +$ python forecast_demo.py +``` +A sample output can be found in the file `ldcast-demo-video-20210622.zip` in the data repository. See the function `forecast_demo` in `forecast_demo.py` see how the `Forecast` class works. To run an ensemble mean of 8 members using the `ForecastDistributed` class, you can use: +```bash +$ python forecast_demo.py --ensemble-members=8 +``` + +The demo for a single ensemble member runs in a couple of minutes on our system using one V100 GPU; with a CPU around 10 minutes or more would be expected. A progress bar will show the status of the generation. + +# Training + +## Training data + +The preprocessed training data, needed to rerun the LDCast training, can be found at the [Zenodo repository](https://doi.org/10.5281/zenodo.7780914). Unzip the `ldcast-datasets.zip` file to the `data` directory. + +## Training the autoencoder + +In the `scripts` directory, run +```bash +$ python train_autoenc.py --model_dir="../models/autoenc_train" +``` +to run the training of the autoencoder with the default parameters. The training checkpoints will be saved in the `../models/autoenc_train` directory (feel free to change this). + +It has been reported that this training may encounter a condition where the loss goes to `nan`. If this happens, try restarting from the latest checkpoint: +```bash +$ python train_autoenc.py --model_dir="../models/autoenc_train" --ckpt_path="../models/autoenc_train/" +``` +where `` should be the latest checkpoint in the `../models/autoenc_train/` directory. + +## Training the diffusion model + +In the `scripts` directory, run +```bash +$ python train_genforecast.py --model_dir="../models/genforecast_train" +``` +to run the training of the diffusion model with the default parameters, or +```bash +$ python train_genforecast.py --model_dir="../models/genforecast_train" --config= +``` +to run the training with different parameters. Some config files can be found in the `config` directory. The training checkpoints will be saved in the `../models/genforecast_train` directory (again, this can be changed freely). + +# Evaluation + +You can find scripts for evaluating models in the `scripts` directory: +* `eval_genforecast.py` to evaluate LDCast +* `eval_dgmr.py` to evaluate DGMR (requires tensorflow installation and the DGMR model from https://github.com/deepmind/deepmind-research/tree/master/nowcasting placed in the `models/dgmr` directory) +* `eval_pysteps.py` to evaluate PySTEPS (requires pysteps installation) +* `metrics.py` to produce metrics from the evaluation results produced with the functions in scripts above +* `plot_genforecast.py` to make plots from the results generated diff --git a/autoenc-32-0.01.pt b/autoenc-32-0.01.pt new file mode 100644 index 0000000000000000000000000000000000000000..3ffb2315591896ebbd8f8a3b8e465463815290ec --- /dev/null +++ b/autoenc-32-0.01.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aa5ad4b8689aadbf702376e7afe5cb437ef5057675e78a8986837e8f28b3126e +size 1617490 diff --git a/config/genforecast-radaronly-128x128-20step.yaml b/config/genforecast-radaronly-128x128-20step.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4534335ded66ab4c403aa51b8402a7e8bea71bab --- /dev/null +++ b/config/genforecast-radaronly-128x128-20step.yaml @@ -0,0 +1 @@ +# this configuration is the default - no parameters to override! diff --git a/config/genforecast-radaronly-256x256-20step.yaml b/config/genforecast-radaronly-256x256-20step.yaml new file mode 100644 index 0000000000000000000000000000000000000000..473db74310984886e163cb8cb313b3b77c29bf98 --- /dev/null +++ b/config/genforecast-radaronly-256x256-20step.yaml @@ -0,0 +1,5 @@ +sample_shape: [8,8] +batch_size: 24 +sampler: "sampler_nowcaster256" +initial_weights: "../models/genforecast/genforecast-radaronly-128x128-20step.pt" +lr: 2.5e-5 \ No newline at end of file diff --git a/environment/environment.yml b/environment/environment.yml new file mode 100644 index 0000000000000000000000000000000000000000..5b63f158e03c0d534c07ae63bb42814d695c6f15 --- /dev/null +++ b/environment/environment.yml @@ -0,0 +1,4 @@ +name: ldcast +channels: + - defaults +prefix: /home/mmhk20/.conda/envs/ldcast diff --git a/environment/ldcast.yml b/environment/ldcast.yml new file mode 100644 index 0000000000000000000000000000000000000000..c06e673ad363f37fcc78e1b622a0e9437019409d --- /dev/null +++ b/environment/ldcast.yml @@ -0,0 +1,170 @@ +name: ldcast +channels: + - conda-forge + - defaults +dependencies: + - _libgcc_mutex=0.1=conda_forge + - _openmp_mutex=4.5=2_gnu + - brotli=1.1.0=hd590300_1 + - brotli-bin=1.1.0=hd590300_1 + - bzip2=1.0.8=hd590300_5 + - c-ares=1.27.0=hd590300_0 + - ca-certificates=2024.2.2=hbcca054_0 + - certifi=2024.2.2=pyhd8ed1ab_0 + - cycler=0.12.1=pyhd8ed1ab_0 + - freetype=2.12.1=h267a509_2 + - geos=3.12.1=h59595ed_0 + - keyutils=1.6.1=h166bdaf_0 + - krb5=1.21.2=h659d440_0 + - lcms2=2.16=hb7c19ff_0 + - ld_impl_linux-64=2.40=h41732ed_0 + - lerc=4.0.0=h27087fc_0 + - libblas=3.9.0=21_linux64_openblas + - libbrotlicommon=1.1.0=hd590300_1 + - libbrotlidec=1.1.0=hd590300_1 + - libbrotlienc=1.1.0=hd590300_1 + - libcblas=3.9.0=21_linux64_openblas + - libcurl=8.6.0=hca28451_0 + - libdeflate=1.19=hd590300_0 + - libedit=3.1.20191231=he28a2e2_2 + - libev=4.33=hd590300_2 + - libexpat=2.6.2=h59595ed_0 + - libffi=3.4.2=h7f98852_5 + - libgcc-ng=13.2.0=h807b86a_5 + - libgfortran-ng=13.2.0=h69a702a_5 + - libgfortran5=13.2.0=ha4646dd_5 + - libgomp=13.2.0=h807b86a_5 + - libjpeg-turbo=3.0.0=hd590300_1 + - liblapack=3.9.0=21_linux64_openblas + - libnghttp2=1.58.0=h47da74e_1 + - libnsl=2.0.1=hd590300_0 + - libopenblas=0.3.26=pthreads_h413a1c8_0 + - libpng=1.6.43=h2797004_0 + - libsqlite=3.45.2=h2797004_0 + - libssh2=1.11.0=h0841786_0 + - libstdcxx-ng=13.2.0=h7e041cc_5 + - libtiff=4.6.0=ha9c0a0a_2 + - libuuid=2.38.1=h0b41bf4_0 + - libwebp-base=1.3.2=hd590300_0 + - libxcb=1.15=h0b41bf4_0 + - libxcrypt=4.4.36=hd590300_1 + - libzlib=1.2.13=hd590300_5 + - matplotlib-base=3.8.3=py312he5832f3_0 + - munkres=1.1.4=pyh9f0ad1d_0 + - ncurses=6.4.20240210=h59595ed_0 + - openjpeg=2.5.2=h488ebb8_0 + - openssl=3.2.1=hd590300_1 + - packaging=24.0=pyhd8ed1ab_0 + - pip=24.0=pyhd8ed1ab_0 + - proj=9.3.1=h1d62c97_0 + - pthread-stubs=0.4=h36c2ea0_1001 + - pyparsing=3.1.2=pyhd8ed1ab_0 + - pyshp=2.3.1=pyhd8ed1ab_0 + - python=3.12.2=hab00c5b_0_cpython + - python-dateutil=2.9.0=pyhd8ed1ab_0 + - python_abi=3.12=4_cp312 + - readline=8.2=h8228510_1 + - setuptools=69.2.0=pyhd8ed1ab_0 + - six=1.16.0=pyh6c4a22f_0 + - sqlite=3.45.2=h2c6b66d_0 + - tk=8.6.13=noxft_h4845f30_101 + - wheel=0.43.0=pyhd8ed1ab_0 + - xorg-libxau=1.0.11=hd590300_0 + - xorg-libxdmcp=1.1.3=h7f98852_0 + - xz=5.2.6=h166bdaf_0 + - zstd=1.5.5=hfc55251_0 + - pip: + - aiobotocore==2.12.1 + - aiohttp==3.9.3 + - aioitertools==0.11.0 + - aiosignal==1.3.1 + - antlr4-python3-runtime==4.9.3 + - arm-pyart==1.18.0 + - attrs==23.2.0 + - botocore==1.34.51 + - cartopy==0.22.0 + - cftime==1.6.3 + - charset-normalizer==3.3.2 + - click==8.1.7 + - cloudpickle==3.0.0 + - cmweather==0.3.2 + - contourpy==1.2.0 + - dask==2024.3.1 + - deprecation==2.1.0 + - einops==0.7.0 + - filelock==3.13.1 + - fire==0.6.0 + - fonttools==4.50.0 + - frozenlist==1.4.1 + - fsspec==2024.3.1 + - h5netcdf==1.3.0 + - h5py==3.10.0 + - idna==3.6 + - jinja2==3.1.3 + - jmespath==1.0.1 + - jsmin==3.0.1 + - jsonschema==4.22.0 + - jsonschema-specifications==2023.12.1 + - kiwisolver==1.4.5 + - lat-lon-parser==1.3.0 + - lightning==2.2.4 + - lightning-utilities==0.11.0 + - llvmlite==0.42.0 + - locket==1.0.0 + - markupsafe==2.1.5 + - matplotlib==3.8.3 + - mda-xdrlib==0.2.0 + - mpmath==1.3.0 + - multidict==6.0.5 + - netcdf4==1.6.5 + - networkx==3.2.1 + - numba==0.59.1 + - numpy==1.26.4 + - nvidia-cublas-cu12==12.1.3.1 + - nvidia-cuda-cupti-cu12==12.1.105 + - nvidia-cuda-nvrtc-cu12==12.1.105 + - nvidia-cuda-runtime-cu12==12.1.105 + - nvidia-cudnn-cu12==8.9.2.26 + - nvidia-cufft-cu12==11.0.2.54 + - nvidia-curand-cu12==10.3.2.106 + - nvidia-cusolver-cu12==11.4.5.107 + - nvidia-cusparse-cu12==12.1.0.106 + - nvidia-nccl-cu12==2.20.5 + - nvidia-nvjitlink-cu12==12.4.99 + - nvidia-nvtx-cu12==12.1.105 + - omegaconf==2.3.0 + - open-radar-data==0.1.0 + - opencv-python==4.9.0.80 + - pandas==2.2.1 + - partd==1.4.1 + - pillow==10.2.0 + - platformdirs==4.2.0 + - pooch==1.8.1 + - pyproj==3.6.1 + - pysteps==1.9.0 + - pytorch-lightning==2.2.1 + - pytz==2024.1 + - pyyaml==6.0.1 + - referencing==0.35.1 + - requests==2.31.0 + - rpds-py==0.18.1 + - s3fs==2024.3.1 + - scipy==1.12.0 + - shapely==2.0.3 + - sympy==1.12 + - termcolor==2.4.0 + - toolz==0.12.1 + - torch==2.3.0 + - torchmetrics==1.3.2 + - torchvision==0.18.0 + - tqdm==4.66.2 + - typing-extensions==4.10.0 + - tzdata==2024.1 + - urllib3==2.0.7 + - wradlib==2.0.3 + - wrapt==1.16.0 + - xarray==2024.2.0 + - xarray-datatree==0.0.14 + - xmltodict==0.13.0 + - xradar==0.4.3 + - yarl==1.9.4 diff --git a/genforecast-radaronly-256x256-20step.pt b/genforecast-radaronly-256x256-20step.pt new file mode 100644 index 0000000000000000000000000000000000000000..e42018092093fabaffbb94ed4e1fd69b95429095 --- /dev/null +++ b/genforecast-radaronly-256x256-20step.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5fef86b78d29fde8ba66f51ae74f0d84ddc67b711fcab034a3130ea5ac7721cf +size 5345469521 diff --git a/ldcast/analysis/confmatrix.py b/ldcast/analysis/confmatrix.py new file mode 100644 index 0000000000000000000000000000000000000000..0c81f39d884b41cbb991ec4071c94a1120c9091b --- /dev/null +++ b/ldcast/analysis/confmatrix.py @@ -0,0 +1,117 @@ +import os +import concurrent +import multiprocessing + +import netCDF4 +import numpy as np +from scipy.integrate import trapezoid + +from ..features.io import load_batch + + +def confusion_matrix(fc_frac, obs_frac, prob_threshold): + N = np.prod(fc_frac.shape) + fc_above = fc_frac > prob_threshold + obs_above = obs_frac > prob_threshold + tp = np.count_nonzero(fc_above & obs_above) / N + fp = np.count_nonzero(fc_above & ~obs_above) / N + fn = np.count_nonzero(~fc_above & obs_above) / N + tn = 1.0 - tp - fp - fn + return np.array(((tp, fn), (fp, tn))) + + +def confusion_matrix_thresholds(fc_frac, obs_frac, thresholds): + N_threads = multiprocessing.cpu_count() + with concurrent.futures.ThreadPoolExecutor(N_threads) as executor: + futures = [ + executor.submit(confusion_matrix, fc_frac, obs_frac, t) + for t in thresholds + ] + conf_matrix = [f.result() for f in futures] + return np.stack(conf_matrix, axis=-1) + + +def confusion_matrix_thresholds_leadtime(fc_frac, obs_frac, thresholds): + N_threads = multiprocessing.cpu_count() + conf_matrix = [] + with concurrent.futures.ThreadPoolExecutor(N_threads) as executor: + for lt in range(fc_frac.shape[2]): + futures = [ + executor.submit(confusion_matrix, + fc_frac[...,lt,:,:], obs_frac[...,lt,:,:], t) + for t in thresholds + ] + conf_matrix_lt = [f.result() for f in futures] + conf_matrix_lt = np.stack(conf_matrix_lt, axis=-1) + conf_matrix.append(conf_matrix_lt) + + return np.stack(conf_matrix, axis=-2) + + + +def precision(conf_matrix): + ((tp, fn), (fp, tn)) = conf_matrix + precision = tp / (tp + fp) + precision[np.isnan(precision)] = 1.0 + return precision + + +def recall(conf_matrix): + ((tp, fn), (fp, tn)) = conf_matrix + return tp / (tp + fn) + + +def false_alarm_ratio(conf_matrix): + return 1.0 - precision(conf_matrix) + + +def intersection_over_union(conf_matrix): + ((tp, fn), (fp, tn)) = conf_matrix + return tp / (tp+fp+fn) + + +def equitable_threat_score(conf_matrix): + ((tp, fn), (fp, tn)) = conf_matrix + tp_rnd = (tp+fn) * (tp+fp) / (tp+fp+tn+fn) + return (tp-tp_rnd) / (tp+fp+fn-tp_rnd) + + +def peirce_skill_score(conf_matrix): + ((tp, fn), (fp, tn)) = conf_matrix + return (tp*tn - fn*fp) / ((tp+fn) * (fp+tn)) + + +def heidke_skill_score(conf_matrix): + ((tp, fn), (fp, tn)) = conf_matrix + return 2 * (tp*tn - fn*fp) / ((tp+fn)*(fn+tn) + (tp+fp)*(fp+tn)) + + +def roc_area_under_curve(conf_matrix): + ((tp, fn), (fp, tn)) = conf_matrix + tpr = tp / (tp + fn) + fpr = fp / (fp + tn) + + auc = trapezoid(tpr[::-1], x=fpr[::-1]) + return auc + + +def pr_area_under_curve(conf_matrix): + prec = precision(conf_matrix) + rec = recall(conf_matrix) + + if (rec[-1] != 0) or (prec[-1] != 1): + rec = np.hstack((rec, 0.0)) + prec = np.hstack((prec, 1.0)) + + auc = trapezoid(prec[::-1], x=rec[::-1]) + return auc + + +def cost_loss_value(conf_matrix, cost, loss, p_clim): + ((tp, fn), (fp, tn)) = conf_matrix + + E_c = min(cost, p_clim*loss) + E_p = p_clim * cost + E_f = (tp+fp)*cost + fn*loss + #print(cost, loss, p_clim, E_c, E_p, E_f[len(E_f)//2]/E_p, (E_f[len(E_f)//2] - E_c) / (E_p - E_c)) + return (E_f - E_c) / (E_p - E_c) diff --git a/ldcast/analysis/crps.py b/ldcast/analysis/crps.py new file mode 100644 index 0000000000000000000000000000000000000000..7e7840b31a44d3c0222f6934f26d2432d1498dc3 --- /dev/null +++ b/ldcast/analysis/crps.py @@ -0,0 +1,162 @@ +import concurrent.futures +import multiprocessing +import os + +import netCDF4 +import numpy as np + +from ..features.io import load_batch, decode_saved_var_to_rainrate + + +def crps_ensemble(observation, forecasts): + shape = observation.shape + N = np.prod(shape) + shape_flat = (np.prod(shape),) + observation = observation.reshape((N,)) + forecasts = forecasts.reshape((N, forecasts.shape[-1])) + crps_all = np.zeros_like(observation) + N_threads = multiprocessing.cpu_count() + + def crps_chunk(k): + i0 = int(round((k/N_threads) * N)) + i1 = int(round(((k+1) / N_threads) * N)) + obs = observation[i0:i1].copy() + fc = forecasts[i0:i1,:].copy() + fc.sort(axis=-1) + fc_below = fc < obs[...,None] + crps = np.zeros_like(obs) + + for i in range(fc.shape[-1]): + below = fc_below[...,i] + weight = ((i+1)**2 - i**2) / fc.shape[-1]**2 + crps[below] += weight * (obs[below]-fc[...,i][below]) + + for i in range(fc.shape[-1]-1,-1,-1): + above = ~fc_below[...,i] + k = fc.shape[-1]-1-i + weight = ((k+1)**2 - k**2) / fc.shape[-1]**2 + crps[above] += weight * (fc[...,i][above]-obs[above]) + + crps_all[i0:i1] = crps + + with concurrent.futures.ThreadPoolExecutor(N_threads) as executor: + futures = {} + for k in range(N_threads): + args = (crps_chunk, k) + futures[executor.submit(*args)] = k + concurrent.futures.wait(futures) + + return crps_all.reshape(shape) + + +def crps_ensemble_multiscale(observation, forecasts): + obs = observation + fc = forecasts + + crps_scales = {} + scale = 1 + while True: + c = crps_ensemble(obs, fc) + crps_scales[scale] = c + scale *= 2 + if obs.shape[-1] == 1: + break + # avg pooling + obs = 0.25 * ( + obs[...,::2,::2] + + obs[...,1::2,::2] + + obs[...,::2,1::2] + + obs[...,1::2,1::2] + ) + fc = 0.25 * ( + fc[...,::2,::2,:] + + fc[...,1::2,::2,:] + + fc[...,::2,1::2,:] + + fc[...,1::2,1::2,:] + ) + + return crps_scales + + +def gather_observation(data_dir): + files = sorted(os.listdir(data_dir)) + files = [os.path.join(data_dir,fn) for fn in files] + + def obs_from_file(fn): + with netCDF4.Dataset(fn, 'r') as ds: + obs = np.array(ds["future_observations"][:], copy=False) + obs = decode_saved_var_to_rainrate(obs) + p = 1 + obs_pooled = {} + while True: + obs_pooled[p] = obs + if obs.shape[-1] == 1: + break + obs = 0.25 * ( + obs[...,::2,::2] + + obs[...,1::2,::2] + + obs[...,::2,1::2] + + obs[...,1::2,1::2] + ) + p *= 2 + return obs_pooled + + obs_pooled = {} + for fn in files: + print(fn) + obs_file = obs_from_file(fn) + for k in obs_file: + if k not in obs_pooled: + obs_pooled[k] = [] + obs_pooled[k].append(obs_file[k]) + + for k in obs_pooled: + obs_pooled[k] = np.concatenate(obs_pooled[k], axis=0) + + return obs_pooled + + +def process_batch(fn, log=False, preproc_fc=None): + print(fn) + (_, y, y_pred) = load_batch(fn, log=log, preproc_fc=preproc_fc) + return crps_ensemble_multiscale(y, y_pred) + + +def save_crps_for_dataset(data_dir, result_fn, log=False, preproc_fc=None): + files = sorted(os.listdir(data_dir)) + files = [os.path.join(data_dir,fn) for fn in files] + + N_threads = multiprocessing.cpu_count() + futures = [] + with concurrent.futures.ProcessPoolExecutor(N_threads) as executor: + for fn in files: + args = (process_batch, fn) + kwargs = {"log": log, "preproc_fc": preproc_fc} + futures.append(executor.submit(*args, **kwargs)) + + crps = [f.result() for f in futures] + scales = sorted(crps[0].keys()) + crps = { + s: np.concatenate([c[s] for c in crps], axis=0) + for s in scales + } + + with netCDF4.Dataset(result_fn, 'w') as ds: + ds.createDimension("dim_sample", crps[1].shape[0]) + ds.createDimension("dim_channel", crps[1].shape[1]) + ds.createDimension("dim_time_future", crps[1].shape[2]) + var_params = {"zlib": True, "complevel": 1} + + for s in scales: + ds.createDimension(f"dim_h_pool{s}x{s}", crps[s].shape[3]) + ds.createDimension(f"dim_w_pool{s}x{s}", crps[s].shape[4]) + var = ds.createVariable( + f"crps_pool{s}x{s}", np.float32, + ( + "dim_sample", "dim_channel", "dim_time_future", + f"dim_h_pool{s}x{s}", f"dim_w_pool{s}x{s}", + ), + chunksizes=(1,)+crps[s].shape[1:], + **var_params + ) + var[:] = crps[s] diff --git a/ldcast/analysis/fss.py b/ldcast/analysis/fss.py new file mode 100644 index 0000000000000000000000000000000000000000..80d984feb90a2c2316423b43b2178f83b88246c2 --- /dev/null +++ b/ldcast/analysis/fss.py @@ -0,0 +1,137 @@ +import concurrent.futures +import multiprocessing +import os + +import netCDF4 +import numpy as np + +from ..features.io import load_batch, decode_saved_var_to_rainrate + + +def fractions_ensemble(observation, forecasts, threshold, max_scale=256): + obs = (observation >= threshold).astype(np.float32) + fc = (forecasts >= threshold).astype(np.float32).mean(axis=-1) + obs_frac = {} + fc_frac = {} + + scale = 1 + while True: + obs_frac[scale] = obs.copy() + fc_frac[scale] = fc.copy() + scale *= 2 + if scale > max_scale: + break + obs = 0.25 * ( + obs[...,::2,::2] + + obs[...,1::2,::2] + + obs[...,::2,1::2] + + obs[...,1::2,1::2] + ) + fc = 0.25 * ( + fc[...,::2,::2] + + fc[...,1::2,::2] + + fc[...,::2,1::2] + + fc[...,1::2,1::2] + ) + + return (obs_frac, fc_frac) + + +def frac_from_file(fn, threshold, preproc_fc): + print(fn) + (_, y, y_pred) = load_batch(fn, preproc_fc=preproc_fc) + return fractions_ensemble(y, y_pred, threshold) + + +def save_fractions_for_dataset(data_dir, result_fn, threshold, preproc_fc=None): + files = sorted(os.listdir(data_dir)) + files = [os.path.join(data_dir,fn) for fn in files] + + N_threads = multiprocessing.cpu_count() + with concurrent.futures.ProcessPoolExecutor(N_threads) as executor: + futures = [] + for fn in files: + args = (frac_from_file, fn, threshold, preproc_fc) + futures.append(executor.submit(*args)) + + (obs_frac, fc_frac) = zip(*(f.result() for f in futures)) + + scales = list(obs_frac[0].keys()) + obs_frac_dict = {} + fc_frac_dict = {} + for s in scales: + obs_frac_dict[s] = np.concatenate([f[s] for f in obs_frac], axis=0) + fc_frac_dict[s] = np.concatenate([f[s] for f in fc_frac], axis=0) + obs_frac = obs_frac_dict + fc_frac = fc_frac_dict + + frac_vars = {} + k = 0 + with netCDF4.Dataset(result_fn, 'w') as ds: + ds.createDimension("dim_sample", obs_frac[1].shape[0]) + ds.createDimension("dim_channel", obs_frac[1].shape[1]) + ds.createDimension("dim_time_future", obs_frac[1].shape[2]) + var_params = {"zlib": True, "complevel": 1} + for s in scales: + ds.createDimension(f"dim_h_pool{s}x{s}", obs_frac[s].shape[3]) + ds.createDimension(f"dim_w_pool{s}x{s}", obs_frac[s].shape[4]) + obs_var = ds.createVariable( + f"obs_frac_scale{s}x{s}", np.float32, + ( + "dim_sample", "dim_channel", "dim_time_future", + f"dim_h_pool{s}x{s}", f"dim_w_pool{s}x{s}", + ), + chunksizes=(1,)+obs_frac[s].shape[1:], + **var_params + ) + obs_var[:] = obs_frac[s] + fc_var = ds.createVariable( + f"fc_frac_scale{s}x{s}", np.float32, + ( + "dim_sample", "dim_channel", "dim_time_future", + f"dim_h_pool{s}x{s}", f"dim_w_pool{s}x{s}", + ), + chunksizes=(1,)+fc_frac[s].shape[1:], + **var_params + ) + fc_var[:] = fc_frac[s] + + +def load_fractions(fn): + obs_frac = {} + fc_frac = {} + with netCDF4.Dataset(fn, 'r') as ds: + var_list = ds.variables.keys() + scales = {int(v.split("x")[-1]) for v in var_list} + for s in scales: + obs_frac[s] = np.array(ds[f"obs_frac_scale{s}x{s}"][:], copy=False) + fc_frac[s] = np.array(ds[f"fc_frac_scale{s}x{s}"][:], copy=False) + + return (obs_frac, fc_frac) + + +def fractions_skill_score( + obs_frac, fc_frac, + frac_axes=None, fss_axes=None, use_timesteps=None +): + if isinstance(obs_frac, dict): + return { + s: fractions_skill_score( + obs_frac[s], fc_frac[s], + frac_axes=frac_axes, fss_axes=fss_axes, + use_timesteps=use_timesteps + ) + for s in sorted(obs_frac) + } + + if use_timesteps != None: + obs_frac = obs_frac[:,:,:use_timesteps,...] + fc_frac = fc_frac[:,:,:use_timesteps,...] + fbs = ((obs_frac - fc_frac)**2).mean(axis=frac_axes) + fbs_ref = (obs_frac**2).mean(axis=frac_axes) + \ + (fc_frac**2).mean(axis=frac_axes) + fss = 1 - fbs/fbs_ref + if isinstance(fss, np.ndarray): + fss[~np.isfinite(fss)] = 1 + fss = fss.mean(axis=fss_axes) + return fss diff --git a/ldcast/analysis/histogram.py b/ldcast/analysis/histogram.py new file mode 100644 index 0000000000000000000000000000000000000000..293757c678feeb27251b8b790153fff84b617a46 --- /dev/null +++ b/ldcast/analysis/histogram.py @@ -0,0 +1,108 @@ +import concurrent.futures +import multiprocessing +import os + +import netCDF4 +import numpy as np +from scipy.interpolate import interp1d + +from ..features.io import load_batch, decode_saved_var_to_rainrate + + +def histogram(observation, forecasts, bins): + N_bins = len(bins)-1 + N_timesteps = observation.shape[2] + obs_hist = np.zeros((N_bins, N_timesteps), dtype=np.uint64) + fc_hist = np.zeros((N_bins, N_timesteps), dtype=np.uint64) + + for t in range(observation.shape[2]): + obs = observation[:,:,t,...].flatten() + fc = forecasts[:,:,t,...].flatten() + obs_hist[:,t] = np.histogram(obs, bins=bins)[0] + fc_hist[:,t] = np.histogram(fc, bins=bins)[0] + + return (obs_hist, fc_hist) + + +def hist_from_file(fn, bins): + print(fn) + (_, y, y_pred) = load_batch(fn, threshold=bins[0]) + return histogram(y, y_pred, bins) + + +def save_histogram_for_dataset(data_dir, result_fn, bins=(0.05,120,100)): + files = sorted(os.listdir(data_dir)) + files = [os.path.join(data_dir,fn) for fn in files] + + bins = np.exp(np.linspace(np.log(bins[0]), np.log(bins[1]), bins[2])) + bins = np.hstack((0, bins)) + + N_threads = multiprocessing.cpu_count() + with concurrent.futures.ProcessPoolExecutor(N_threads) as executor: + futures = [] + for fn in files: + args = (hist_from_file, fn, bins) + futures.append(executor.submit(*args)) + + (obs_hist, fc_hist) = zip(*(f.result() for f in futures)) + + obs_hist = sum(obs_hist) + fc_hist = sum(fc_hist) + + with netCDF4.Dataset(result_fn, 'w') as ds: + ds.createDimension("dim_bin", obs_hist.shape[0]) + ds.createDimension("dim_time_future", obs_hist.shape[1]) + var_params = {"zlib": True, "complevel": 1} + + obs_var = ds.createVariable( + f"obs_hist", np.uint64, + ("dim_bin", "dim_time_future"), + **var_params + ) + obs_var[:] = obs_hist + + fc_var = ds.createVariable( + f"fc_hist", np.uint64, + ("dim_bin", "dim_time_future"), + **var_params + ) + fc_var[:] = fc_hist + + ds.createDimension("dim_bin_edge", len(bins)) + bin_var = ds.createVariable( + f"bins", np.float64, + ("dim_bin_edge",), + **var_params + ) + bin_var[:] = bins + + +def load_histogram(fn): + with netCDF4.Dataset(fn, 'r') as ds: + obs_hist = np.array(ds["obs_hist"][:], copy=False) + fc_hist = np.array(ds["fc_hist"][:], copy=False) + bins = np.array(ds["bins"][:], copy=False) + + return (obs_hist, fc_hist, bins) + + +class ProbabilityMatch: + def __init__(self, obs_hist, fc_hist, bins): + obs_c = obs_hist.cumsum() + obs_c = obs_c / obs_c[-1] + fc_c = fc_hist.cumsum() + fc_c = fc_c / fc_c[-1] + + self.obs_cdf = interp1d(np.hstack((0,obs_c)), bins, fill_value='extrapolate') + self.fc_cdf = interp1d(bins, np.hstack((0,fc_c)), fill_value='extrapolate') + + def __call__(self, x): + return self.obs_cdf(self.fc_cdf(x)) + + +def probability_match_timesteps(obs_hist, fc_hist, bins): + num_timesteps = obs_hist.shape[1] + return [ + ProbabilityMatch(obs_hist[:,t], fc_hist[:,t], bins) + for t in range(num_timesteps) + ] diff --git a/ldcast/analysis/rank.py b/ldcast/analysis/rank.py new file mode 100644 index 0000000000000000000000000000000000000000..df95ce749e0d45a81e6d5e317c428a9c55b8e6fd --- /dev/null +++ b/ldcast/analysis/rank.py @@ -0,0 +1,190 @@ +import os +import concurrent +import multiprocessing + +import netCDF4 +import numpy as np + +from ..features.io import load_batch + + +def ranks_ensemble( + observation, forecasts, + noise_scale=1e-6, rng=None +): + shape = observation.shape + N = np.prod(shape) + shape_flat = (np.prod(shape),) + observation = observation.reshape((N,)) + forecasts = forecasts.reshape((N, forecasts.shape[-1])) + N_threads = multiprocessing.cpu_count() + + max_rank = forecasts.shape[-1] + bins = np.arange(-0.5, max_rank+0.6) + ranks_all = np.zeros_like(observation, dtype=np.uint32) + + if rng is None: + rng = np.random + + def rank_dist_chunk(k): + i0 = int(round((k/N_threads) * N)) + i1 = int(round(((k+1) / N_threads) * N)) + obs = observation[i0:i1].astype(np.float64, copy=True) + fc = forecasts[i0:i1,:].astype(np.float64, copy=True) + + # add a tiny amount of noise to forecast to randomize ties + # (important to add to both obs and fc!) + obs += (rng.rand(*obs.shape) - 0.5) * noise_scale + fc += (rng.rand(*fc.shape) - 0.5) * noise_scale + + ranks = np.count_nonzero(obs[...,None] >= fc, axis=-1) + ranks_all[i0:i1] = ranks + + with concurrent.futures.ThreadPoolExecutor(N_threads) as executor: + futures = {} + for k in range(N_threads): + args = (rank_dist_chunk, k) + futures[executor.submit(*args)] = k + concurrent.futures.wait(futures) + + return ranks_all.reshape(shape) + + +def ranks_multiscale(observation, forecasts): + obs = observation + fc = forecasts + + rank_scales = {} + scale = 1 + while True: + r = ranks_ensemble(obs, fc) + rank_scales[scale] = r + scale *= 2 + if obs.shape[-1] == 1: + break + # avg pooling + obs = 0.25 * ( + obs[...,::2,::2] + + obs[...,1::2,::2] + + obs[...,::2,1::2] + + obs[...,1::2,1::2] + ) + fc = 0.25 * ( + fc[...,::2,::2,:] + + fc[...,1::2,::2,:] + + fc[...,::2,1::2,:] + + fc[...,1::2,1::2,:] + ) + + return rank_scales + + +def rank_distribution(ranks, num_forecasts=32): + N = np.prod(ranks.shape) + bins = np.arange(-0.5, num_forecasts+0.6) + N_threads = multiprocessing.cpu_count() + ranks = ranks.ravel() + + hist = [None] * N_threads + def hist_chunk(k): + i0 = int(round((k/N_threads) * N)) + i1 = int(round(((k+1) / N_threads) * N)) + (h, _) = np.histogram(ranks[i0:i1], bins=bins) + hist[k] = h + + with concurrent.futures.ThreadPoolExecutor(N_threads) as executor: + futures = {} + for k in range(N_threads): + args = (hist_chunk, k) + futures[executor.submit(*args)] = k + concurrent.futures.wait(futures) + + hist = sum(hist) + return hist / hist.sum() + + +def rank_KS(rank_dist, num_forecasts=32): + h = rank_dist + h = h / h.sum() + ch = np.cumsum(h) + cb = np.linspace(0, 1, len(ch)) + return abs(ch-cb).max() + + +def rank_DKL(rank_dist, num_forecasts=32): + h = rank_dist + q = h / h.sum() + p = 1/len(h) + return p*np.log(p/q).sum() + + +def rank_metric_by_leadtime(ranks, metric=None, num_forecasts=32): + if metric is None: + metric = rank_DKL + + metric_by_leadtime = [] + for t in range(ranks.shape[2]): + ranks_time = ranks[:,:,t,...] + h = rank_distribution(ranks_time) + m = metric(h, num_forecasts=num_forecasts) + metric_by_leadtime.append(m) + return np.array(metric_by_leadtime) + + +def rank_metric_by_bin(ranks, values, bins, metric=None, num_forecasts=32): + if metric is None: + metric = rank_DKL + + metric_by_bin = [] + for (b0,b1) in zip(bins[:-1],bins[1:]): + ranks_bin = ranks[(b0 <= values) & (values < b1)] + h = rank_distribution(ranks_bin) + m = metric(h, num_forecasts=num_forecasts) + metric_by_bin.append(m) + return np.array(metric_by_bin) + + +def process_batch(fn, preproc_fc=None): + print(fn) + (_, y, y_pred) = load_batch(fn, preproc_fc=preproc_fc) + return ranks_multiscale(y, y_pred) + + +def save_ranks_for_dataset(data_dir, result_fn, preproc_fc=None): + files = sorted(os.listdir(data_dir)) + files = [os.path.join(data_dir,fn) for fn in files] + + N_threads = multiprocessing.cpu_count() + futures = [] + with concurrent.futures.ProcessPoolExecutor(N_threads) as executor: + for fn in files: + args = (process_batch, fn) + kwargs = {"preproc_fc": preproc_fc} + futures.append(executor.submit(*args, **kwargs)) + + ranks = [f.result() for f in futures] + scales = sorted(ranks[0].keys()) + ranks = { + s: np.concatenate([r[s] for r in ranks], axis=0) + for s in scales + } + + with netCDF4.Dataset(result_fn, 'w') as ds: + ds.createDimension("dim_sample", ranks[1].shape[0]) + ds.createDimension("dim_channel", ranks[1].shape[1]) + ds.createDimension("dim_time_future", ranks[1].shape[2]) + var_params = {"zlib": True, "complevel": 1} + + for s in scales: + ds.createDimension(f"dim_h_pool{s}x{s}", ranks[s].shape[3]) + ds.createDimension(f"dim_w_pool{s}x{s}", ranks[s].shape[4]) + var = ds.createVariable( + f"ranks_pool{s}x{s}", np.float32, + ( + "dim_sample", "dim_channel", "dim_time_future", + f"dim_h_pool{s}x{s}", f"dim_w_pool{s}x{s}", + ), + chunksizes=(1,)+ranks[s].shape[1:], + **var_params + ) + var[:] = ranks[s] diff --git a/ldcast/features/.sampling.py.swp b/ldcast/features/.sampling.py.swp new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/ldcast/features/batch.py b/ldcast/features/batch.py new file mode 100644 index 0000000000000000000000000000000000000000..00cc38cdd349a9dc222ad5438c630521d148ed9e --- /dev/null +++ b/ldcast/features/batch.py @@ -0,0 +1,375 @@ +from datetime import datetime, timedelta +import os +import pickle + +from numba import njit, prange, types +from numba.typed import Dict +import numpy as np +from torch.utils.data import Dataset, IterableDataset + +from .patches import unpack_patches +from .sampling import EqualFrequencySampler + + +class BatchGenerator: + def __init__(self, + variables, + raw, + predictors, + target, + primary_var, + time_range_sampling=(-1,2), + forecast_raw_vars=(), + sampling_bins=None, + sampler_file=None, + sample_shape=(4,4), + batch_size=32, + interval=timedelta(minutes=5), + random_seed=None, + augment=False + ): + super().__init__() + self.batch_size = batch_size + self.interval = interval + self.interval_secs = np.int64(self.interval.total_seconds()) + self.variables = variables + self.predictors = predictors + self.target = target + self.used_variables = predictors + [target] + self.rng = np.random.RandomState(seed=random_seed) + self.augment = augment + + # setup indices for retrieving source raw data + self.sources = set.union( + *(set(variables[v]["sources"]) for v in self.used_variables) + ) + self.forecast_raw_vars = set(forecast_raw_vars) & self.sources + self.patch_index = {} + for raw_name_base in self.sources: + if raw_name_base in forecast_raw_vars: + raw_names = ( + rn for rn in raw if rn.startswith(raw_name_base+"-") + ) + else: + raw_names = (raw_name_base,) + for raw_name in raw_names: + raw_data = raw[raw_name] + self.setup_index(raw_name, raw_data, sample_shape) + + for raw_name in self.forecast_raw_vars: + patch_index_var = { + k: v for (k,v) in self.patch_index.items() + if k.startswith(raw_name+"-") + } + self.patch_index[raw_name] = \ + ForecastPatchIndexWrapper(patch_index_var) + + # setup samplers + if (sampler_file is None) or not os.path.isfile(sampler_file): + print("No cached sampler found, creating a new one...") + primary_raw_var = variables[primary_var]["sources"][0] + t0 = t1 = None + for (var_name, var_data) in variables.items(): + timesteps = var_data["timesteps"][[0,-1]].copy() + timesteps[0] -= 1 + ts_secs = timesteps * \ + var_data.get("timestep_secs", self.interval_secs) + timesteps = ts_secs // self.interval_secs + t0 = timesteps[0] if t0 is None else min(t0,timesteps[0]) + t1 = timesteps[-1] if t1 is None else max(t1,timesteps[-1]) + time_range_valid = (t0,t1+1) + self.sampler = EqualFrequencySampler( + sampling_bins, raw[primary_raw_var], + self.patch_index[primary_raw_var], sample_shape, + time_range_valid, time_range_sampling=time_range_sampling, + timestep_secs=self.interval_secs + ) + if sampler_file is not None: + print(f"Caching sampler to {sampler_file}.") + with open(sampler_file, 'wb') as f: + pickle.dump(self.sampler, f) + else: + print(f"Loading cached sampler from {sampler_file}.") + with open(sampler_file, 'rb') as f: + self.sampler = pickle.load(f) + + def setup_index(self, raw_name, raw_data, box_size): + zero_value = raw_data.get("zero_value", 0) + missing_value = raw_data.get("missing_value", zero_value) + + self.patch_index[raw_name] = PatchIndex( + *unpack_patches(raw_data), + zero_value=zero_value, + missing_value=missing_value, + interval=self.interval, + box_size=box_size + ) + + def augmentations(self): + return tuple(self.rng.randint(2, size=3)) + + def augment_batch(self, batch, transpose, flipud, fliplr): + if self.augment: + if transpose: + axes = list(range(batch.ndim)) + axes = axes[:-2] + [axes[-1], axes[-2]] + batch = batch.transpose(axes) + flips = [] + if flipud: + flips.append(-2) + if fliplr: + flips.append(-1) + if flips: + batch = np.flip(batch, axis=flips) + return batch.copy() + + def batch(self, samples=None, batch_size=None): + if batch_size is None: + batch_size = self.batch_size + + if samples is None: + # get the sample coordinates from the sampler + samples = self.sampler(batch_size) + + print(samples) + (t0,i0,j0) = samples.T + + if self.augment: + augmentations = self.augmentations() + + batch = {} + for var_name in self.used_variables: + var_data = self.variables[var_name] + + # different timestep from standard (e.g. forecast); round down + # to times where we have data available + ts_secs = var_data.get("timestep_secs", self.interval_secs) + t_shift = -(t0 % ts_secs) + t0_shifted = t0 + t_shift + t = t0_shifted[:,None] + ts_secs*var_data["timesteps"][None,:] + t_relative = (t - t0[:,None]) / self.interval_secs + + # read raw data from index + raw_data = ( + self.patch_index[raw_name](t,i0,j0) + for raw_name in var_data["sources"] + ) + + # transform to model variable + batch_var = var_data["transform"](*raw_data) + + # add channel dimension if not already present + add_dims = (1,) if batch_var.ndim == 4 else () + batch_var = np.expand_dims(batch_var, add_dims) + + # data augmentation + if self.augment: + batch_var = self.augment_batch(batch_var, *augmentations) + + # bundle with time coordinates + batch[var_name] = (batch_var, t_relative.astype(np.float32)) + + pred_batch = [batch[v] for v in self.predictors] + target_batch = batch[self.target][0] # no time coordinates for target + return (pred_batch, target_batch) + + def batches(self, *args, num=None, **kwargs): + if num is not None: + for i in range(num): + yield self.batch(*args, **kwargs) + else: + while True: + yield self.batch(*args, **kwargs) + + +class StreamBatchDataset(IterableDataset): + def __init__(self, batch_gen, batches_per_epoch): + super().__init__() + self.batch_gen = batch_gen + self.batches_per_epoch = batches_per_epoch + + def __iter__(self): + batches = self.batch_gen.batches(num=self.batches_per_epoch) + yield from batches + + +class DeterministicBatchDataset(Dataset): + def __init__(self, batch_gen, batches_per_epoch, random_seed=None): + super().__init__() + self.batch_gen = batch_gen + self.batches_per_epoch = batches_per_epoch + self.batch_gen.sampler.rng = np.random.RandomState(seed=random_seed) + self.samples = [ + self.batch_gen.sampler(self.batch_gen.batch_size) + for i in range(self.batches_per_epoch) + ] + + def __len__(self): + return self.batches_per_epoch + + def __getitem__(self, ind): + print(self.samples[ind]) + return self.batch_gen.batch(samples=self.samples[ind]) + + +class PatchIndex: + IDX_ZERO = -1 + IDX_MISSING = -2 + + def __init__( + self, patch_data, patch_coords, patch_times, + zero_patch_coords, zero_patch_times, + interval=timedelta(minutes=5), + box_size=(4,4), zero_value=0, + missing_value=0 + ): + self.dt = int(round(interval.total_seconds())) + self.box_size = box_size + self.zero_value = zero_value + self.missing_value = missing_value + self.patch_data = patch_data + self.sample_shape = ( + self.patch_data.shape[1]*box_size[0], + self.patch_data.shape[2]*box_size[1] + ) + + self.patch_index = Dict.empty( + key_type=types.UniTuple(types.int64, 3), + value_type=types.int64 + ) + init_patch_index(self.patch_index, patch_coords, patch_times) + init_patch_index_zero(self.patch_index, zero_patch_coords, + zero_patch_times, PatchIndex.IDX_ZERO) + + self._batch = None + + def _alloc_batch(self, batch_size, num_timesteps): + needs_rebuild = (self._batch is None) or \ + (self._batch.shape[0] < batch_size) or \ + (self._batch.shape[1] < num_timesteps) + if needs_rebuild: + del self._batch + self._batch = np.zeros( + (batch_size,num_timesteps)+self.sample_shape, + self.patch_data.dtype + ) + return self._batch + + def __call__(self, t, i0_all, j0_all): + batch = self._alloc_batch(*t.shape) + + i1_all = i0_all + self.box_size[0] + j1_all = j0_all + self.box_size[1] + bi_size = self.patch_data.shape[1] + bj_size = self.patch_data.shape[2] + + build_batch(batch, self.patch_data, self.patch_index, + t, i0_all, i1_all, j0_all, j1_all, + bi_size, bj_size, self.zero_value, + self.missing_value) + + return batch[:,:t.shape[1],...] + + +@njit +def init_patch_index(patch_index, patch_coords, patch_times): + for k in range(patch_coords.shape[0]): + t = patch_times[k] + i = np.int64(patch_coords[k,0]) + j = np.int64(patch_coords[k,1]) + patch_index[(t,i,j)] = k + + +@njit +def init_patch_index_zero(patch_index, zero_patch_coords, + zero_patch_times, idx_zero): + + for k in range(zero_patch_coords.shape[0]): + t = zero_patch_times[k] + i = np.int64(zero_patch_coords[k,0]) + j = np.int64(zero_patch_coords[k,1]) + patch_index[(t,i,j)] = idx_zero + + +# numba can't find these values from PatchIndex +IDX_ZERO = PatchIndex.IDX_ZERO +IDX_MISSING = PatchIndex.IDX_MISSING +@njit(parallel=True) +def build_batch( + batch, patch_data, patch_index, + t_all, i0_all, i1_all, j0_all, j1_all, + bi_size, bj_size, zero_value, missing_value +): + for k in prange(t_all.shape[0]): + i0 = i0_all[k] + i1 = i1_all[k] + j0 = j0_all[k] + j1 = j1_all[k] + + for (bt,t) in enumerate(t_all[k,:]): + for i in range(i0, i1): + bi0 = (i-i0) * bi_size + bi1 = bi0 + bi_size + for j in range(j0, j1): + ind = int(patch_index.get((t,i,j), IDX_MISSING)) + bj0 = (j-j0) * bj_size + bj1 = bj0 + bj_size + if ind >= 0: + batch[k,bt,bi0:bi1,bj0:bj1] = patch_data[ind] + elif ind == IDX_ZERO: + batch[k,bt,bi0:bi1,bj0:bj1] = zero_value + elif ind == IDX_MISSING: + batch[k,bt,bi0:bi1,bj0:bj1] = missing_value + + +class ForecastPatchIndexWrapper(PatchIndex): + def __init__(self, patch_index): + self.patch_index = patch_index + raw_names = {"-".join(v.split("-")[:-1]) for v in patch_index} + if len(raw_names) != 1: + raise ValueError( + "Can only wrap variables with the same base name") + self.raw_name = list(raw_names)[0] + lags_hour = [int(v.split("-")[-1]) for v in patch_index] + self.lags_hour = set(lags_hour) + forecast_interval_hour = np.diff(sorted(lags_hour)) + if len(set(forecast_interval_hour)) != 1: + raise ValueError("Lags must be evenly spaced") + forecast_interval_hour = forecast_interval_hour[0] + if (24 % forecast_interval_hour): + raise ValueError( + "24 hours must be a multiple of the forecast interval") + self.forecast_interval_hour = forecast_interval_hour + self.forecast_interval = 3600 * forecast_interval_hour + + # need to set these for _alloc_batch to work + self._batch = None + v = list(self.patch_index.keys())[0] + self.sample_shape = self.patch_index[v].sample_shape + self.patch_data = self.patch_index[v].patch_data + + def __call__(self, t, i0, j0): + batch = self._alloc_batch(*t.shape) + + # ensure that all data come from the same forecast + t0 = t[:,:1] + start_time_from_fc = t0 % self.forecast_interval + time_from_fc = start_time_from_fc + (t - t0) + lags_hour = (time_from_fc // self.forecast_interval) * \ + self.forecast_interval_hour + + for lag in self.lags_hour: + raw_name_lag = f"{self.raw_name}-{lag}" + batch_lag = self.patch_index[raw_name_lag](t,i0,j0) + lag_mask = (lags_hour == lag) + copy_masked_times(batch_lag, batch, lag_mask) + + return batch[:,:t.shape[1],...] + + +@njit(parallel=True) +def copy_masked_times(from_batch, to_batch, mask): + for k in prange(from_batch.shape[0]): + for bt in range(from_batch.shape[1]): + if mask[k,bt]: + to_batch[k,bt,:,:] = from_batch[k,bt,:,:] diff --git a/ldcast/features/batch.py.save b/ldcast/features/batch.py.save new file mode 100644 index 0000000000000000000000000000000000000000..2001e31681f75fff51a70e2cd174139bba9b97cb --- /dev/null +++ b/ldcast/features/batch.py.save @@ -0,0 +1,378 @@ +from datetime import datetime, timedelta +import os +import pickle + +from numba import njit, prange, types +from numba.typed import Dict +import numpy as np +from torch.utils.data import Dataset, IterableDataset + +from .patches import unpack_patches +from .sampling import EqualFrequencySampler + + +class BatchGenerator: + def __init__(self, + variables, + raw, + predictors, + target, + primary_var, + time_range_sampling=(-1,2), + forecast_raw_vars=(), + sampling_bins=None, + sampler_file=None, + sample_shape=(4,4), + batch_size=32, + interval=timedelta(minutes=5), + random_seed=None, + augment=False + ): + super().__init__() + self.batch_size = batch_size + self.interval = interval + self.interval_secs = np.int64(self.interval.total_seconds()) + self.variables = variables + self.predictors = predictors + self.target = target + self.used_variables = predictors + [target] + self.rng = np.random.RandomState(seed=random_seed) + self.augment = augment + + # setup indices for retrieving source raw data + self.sources = set.union( + *(set(variables[v]["sources"]) for v in self.used_variables) + ) + self.forecast_raw_vars = set(forecast_raw_vars) & self.sources + self.patch_index = {} + for raw_name_base in self.sources: + if raw_name_base in forecast_raw_vars: + raw_names = ( + rn for rn in raw if rn.startswith(raw_name_base+"-") + ) + else: + raw_names = (raw_name_base,) + for raw_name in raw_names: + raw_data = raw[raw_name] + self.setup_index(raw_name, raw_data, sample_shape) + + for raw_name in self.forecast_raw_vars: + patch_index_var = { + k: v for (k,v) in self.patch_index.items() + if k.startswith(raw_name+"-") + } + self.patch_index[raw_name] = \ + ForecastPatchIndexWrapper(patch_index_var) + + # setup samplers + if (sampler_file is None) or not os.path.isfile(sampler_file): + print("No cached sampler found, creating a new one...") + primary_raw_var = variables[primary_var]["sources"][0] + t0 = t1 = None + for (var_name, var_data) in variables.items(): + timesteps = var_data["timesteps"][[0,-1]].copy() + timesteps[0] -= 1 + ts_secs = timesteps * \ + var_data.get("timestep_secs", self.interval_secs) + timesteps = ts_secs // self.interval_secs + t0 = timesteps[0] if t0 is None else min(t0,timesteps[0]) + t1 = timesteps[-1] if t1 is None else max(t1,timesteps[-1]) + time_range_valid = (t0,t1+1) + self.sampler = EqualFrequencySampler( + sampling_bins, raw[primary_raw_var], + self.patch_index[primary_raw_var], sample_shape, + time_range_valid, time_range_sampling=time_range_sampling, + timestep_secs=self.interval_secs + ) + if sampler_file is not None: + print(f"Caching sampler to {sampler_file}.") + with open(sampler_file, 'wb') as f: + pickle.dump(self.sampler, f) + else: + print(f"Loading cached sampler from {sampler_file}.") + with open(sampler_file, 'rb') as f: + self.sampler = pickle.load(f) + + def setup_index(self, raw_name, raw_data, box_size): + zero_value = raw_data.get("zero_value", 0) + missing_value = raw_data.get("missing_value", zero_value) + + self.patch_index[raw_name] = PatchIndex( + *unpack_patches(raw_data), + zero_value=zero_value, + missing_value=missing_value, + interval=self.interval, + box_size=box_size + ) + + def augmentations(self): + return tuple(self.rng.randint(2, size=3)) + + def augment_batch(self, batch, transpose, flipud, fliplr): + if self.augment: + if transpose: + axes = list(range(batch.ndim)) + axes = axes[:-2] + [axes[-1], axes[-2]] + batch = batch.transpose(axes) + flips = [] + if flipud: + flips.append(-2) + if fliplr: + flips.append(-1) + if flips: + batch = np.flip(batch, axis=flips) + return batch.copy() + + def batch(self, samples=None, batch_size=None): + if batch_size is None: + batch_size = self.batch_size + + if samples is None: + # get the sample coordinates from the sampler + samples = self.sampler(batch_size) + + print(samples) + (t0,i0,j0) = samples.T + + if self.augment: + augmentations = self.augmentations() + + batch = {} + + + + for var_name in self.used_variables: + var_data = self.variables[var_name] + + # different timestep from standard (e.g. forecast); round down + # to times where we have data available + ts_secs = var_data.get("timestep_secs", self.interval_secs) + t_shift = -(t0 % ts_secs) + t0_shifted = t0 + t_shift + t = t0_shifted[:,None] + ts_secs*var_data["timesteps"][None,:] + t_relative = (t - t0[:,None]) / self.interval_secs + + # read raw data from index + raw_data = ( + self.patch_index[raw_name](t,i0,j0) + for raw_name in var_data["sources"] + ) + + # transform to model variable + batch_var = var_data["transform"](*raw_data) + + # add channel dimension if not already present + add_dims = (1,) if batch_var.ndim == 4 else () + batch_var = np.expand_dims(batch_var, add_dims) + + # data augmentation + if self.augment: + batch_var = self.augment_batch(batch_var, *augmentations) + + # bundle with time coordinates + batch[var_name] = (batch_var, t_relative.astype(np.float32)) + + pred_batch = [batch[v] for v in self.predictors] + target_batch = batch[self.target][0] # no time coordinates for target + return (pred_batch, target_batch) + + def batches(self, *args, num=None, **kwargs): + if num is not None: + for i in range(num): + yield self.batch(*args, **kwargs) + else: + while True: + yield self.batch(*args, **kwargs) + + +class StreamBatchDataset(IterableDataset): + def __init__(self, batch_gen, batches_per_epoch): + super().__init__() + self.batch_gen = batch_gen + self.batches_per_epoch = batches_per_epoch + + def __iter__(self): + batches = self.batch_gen.batches(num=self.batches_per_epoch) + yield from batches + + +class DeterministicBatchDataset(Dataset): + def __init__(self, batch_gen, batches_per_epoch, random_seed=None): + super().__init__() + self.batch_gen = batch_gen + self.batches_per_epoch = batches_per_epoch + self.batch_gen.sampler.rng = np.random.RandomState(seed=random_seed) + self.samples = [ + self.batch_gen.sampler(self.batch_gen.batch_size) + for i in range(self.batches_per_epoch) + ] + + def __len__(self): + return self.batches_per_epoch + + def __getitem__(self, ind): + print(self.samples[ind]) + return self.batch_gen.batch(samples=self.samples[ind]) + + +class PatchIndex: + IDX_ZERO = -1 + IDX_MISSING = -2 + + def __init__( + self, patch_data, patch_coords, patch_times, + zero_patch_coords, zero_patch_times, + interval=timedelta(minutes=5), + box_size=(4,4), zero_value=0, + missing_value=0 + ): + self.dt = int(round(interval.total_seconds())) + self.box_size = box_size + self.zero_value = zero_value + self.missing_value = missing_value + self.patch_data = patch_data + self.sample_shape = ( + self.patch_data.shape[1]*box_size[0], + self.patch_data.shape[2]*box_size[1] + ) + + self.patch_index = Dict.empty( + key_type=types.UniTuple(types.int64, 3), + value_type=types.int64 + ) + init_patch_index(self.patch_index, patch_coords, patch_times) + init_patch_index_zero(self.patch_index, zero_patch_coords, + zero_patch_times, PatchIndex.IDX_ZERO) + + self._batch = None + + def _alloc_batch(self, batch_size, num_timesteps): + needs_rebuild = (self._batch is None) or \ + (self._batch.shape[0] < batch_size) or \ + (self._batch.shape[1] < num_timesteps) + if needs_rebuild: + del self._batch + self._batch = np.zeros( + (batch_size,num_timesteps)+self.sample_shape, + self.patch_data.dtype + ) + return self._batch + + def __call__(self, t, i0_all, j0_all): + batch = self._alloc_batch(*t.shape) + + i1_all = i0_all + self.box_size[0] + j1_all = j0_all + self.box_size[1] + bi_size = self.patch_data.shape[1] + bj_size = self.patch_data.shape[2] + + build_batch(batch, self.patch_data, self.patch_index, + t, i0_all, i1_all, j0_all, j1_all, + bi_size, bj_size, self.zero_value, + self.missing_value) + + return batch[:,:t.shape[1],...] + + +@njit +def init_patch_index(patch_index, patch_coords, patch_times): + for k in range(patch_coords.shape[0]): + t = patch_times[k] + i = np.int64(patch_coords[k,0]) + j = np.int64(patch_coords[k,1]) + patch_index[(t,i,j)] = k + + +@njit +def init_patch_index_zero(patch_index, zero_patch_coords, + zero_patch_times, idx_zero): + + for k in range(zero_patch_coords.shape[0]): + t = zero_patch_times[k] + i = np.int64(zero_patch_coords[k,0]) + j = np.int64(zero_patch_coords[k,1]) + patch_index[(t,i,j)] = idx_zero + + +# numba can't find these values from PatchIndex +IDX_ZERO = PatchIndex.IDX_ZERO +IDX_MISSING = PatchIndex.IDX_MISSING +@njit(parallel=True) +def build_batch( + batch, patch_data, patch_index, + t_all, i0_all, i1_all, j0_all, j1_all, + bi_size, bj_size, zero_value, missing_value +): + for k in prange(t_all.shape[0]): + i0 = i0_all[k] + i1 = i1_all[k] + j0 = j0_all[k] + j1 = j1_all[k] + + for (bt,t) in enumerate(t_all[k,:]): + for i in range(i0, i1): + bi0 = (i-i0) * bi_size + bi1 = bi0 + bi_size + for j in range(j0, j1): + ind = int(patch_index.get((t,i,j), IDX_MISSING)) + bj0 = (j-j0) * bj_size + bj1 = bj0 + bj_size + if ind >= 0: + batch[k,bt,bi0:bi1,bj0:bj1] = patch_data[ind] + elif ind == IDX_ZERO: + batch[k,bt,bi0:bi1,bj0:bj1] = zero_value + elif ind == IDX_MISSING: + batch[k,bt,bi0:bi1,bj0:bj1] = missing_value + + +class ForecastPatchIndexWrapper(PatchIndex): + def __init__(self, patch_index): + self.patch_index = patch_index + raw_names = {"-".join(v.split("-")[:-1]) for v in patch_index} + if len(raw_names) != 1: + raise ValueError( + "Can only wrap variables with the same base name") + self.raw_name = list(raw_names)[0] + lags_hour = [int(v.split("-")[-1]) for v in patch_index] + self.lags_hour = set(lags_hour) + forecast_interval_hour = np.diff(sorted(lags_hour)) + if len(set(forecast_interval_hour)) != 1: + raise ValueError("Lags must be evenly spaced") + forecast_interval_hour = forecast_interval_hour[0] + if (24 % forecast_interval_hour): + raise ValueError( + "24 hours must be a multiple of the forecast interval") + self.forecast_interval_hour = forecast_interval_hour + self.forecast_interval = 3600 * forecast_interval_hour + + # need to set these for _alloc_batch to work + self._batch = None + v = list(self.patch_index.keys())[0] + self.sample_shape = self.patch_index[v].sample_shape + self.patch_data = self.patch_index[v].patch_data + + def __call__(self, t, i0, j0): + batch = self._alloc_batch(*t.shape) + + # ensure that all data come from the same forecast + t0 = t[:,:1] + start_time_from_fc = t0 % self.forecast_interval + time_from_fc = start_time_from_fc + (t - t0) + lags_hour = (time_from_fc // self.forecast_interval) * \ + self.forecast_interval_hour + + for lag in self.lags_hour: + raw_name_lag = f"{self.raw_name}-{lag}" + batch_lag = self.patch_index[raw_name_lag](t,i0,j0) + lag_mask = (lags_hour == lag) + copy_masked_times(batch_lag, batch, lag_mask) + + return batch[:,:t.shape[1],...] + + +@njit(parallel=True) +def copy_masked_times(from_batch, to_batch, mask): + for k in prange(from_batch.shape[0]): + for bt in range(from_batch.shape[1]): + if mask[k,bt]: + to_batch[k,bt,:,:] = from_batch[k,bt,:,:] diff --git a/ldcast/features/io.py b/ldcast/features/io.py new file mode 100644 index 0000000000000000000000000000000000000000..845faba9d7b21bb6b751ba3585c297f88ed605b9 --- /dev/null +++ b/ldcast/features/io.py @@ -0,0 +1,125 @@ +import os + +import netCDF4 +import numpy as np + + +def convert_var_for_saving( + x, fill_value=0.02, min_value=0.05, max_value=118.428, + mean=-0.051, std=0.528 +): + y = x*std + mean + log_min = np.log10(min_value) + log_max = np.log10(max_value) + mask = (y >= log_min) + y = y[mask].clip(max=log_max) + y = (y-log_min) / (log_max-log_min) + yc = np.zeros_like(x, dtype=np.uint16) + yc[mask] = (y*65533).round().astype(np.uint16) + 1 + return yc + + +def decode_saved_var_to_rainrate( + x, fill_value=0.02, min_value=0.05, threshold=0.1, max_value=118.428, + mean=-0.051, std=0.528, log=False, preproc=None +): + mask = (x >= 1) + log_min = np.log10(min_value) + log_max = np.log10(max_value) + yc = log_min + (x[mask].astype(np.float32)-1) * \ + ((log_max-log_min) / 65533) + y = np.zeros_like(x, dtype=np.float32) + + yc = 10**yc + y[mask] = yc + if preproc is not None: + y = [preproc[t](y[:,:,t,...]) for t in range(y.shape[2])] + y = np.stack(y, axis=2) + + if log: + y[y < threshold] = fill_value + y = np.log10(y) + else: + y[y < threshold] = 0.0 + + return y + + +def save_batch(x, y, y_pred, batch_index, fn_template, out_dir, out_fn=None): + while isinstance(x, list) or isinstance(x, tuple): + x = x[0] + + x = convert_var_for_saving(np.array(x, copy=False)) + y = convert_var_for_saving(np.array(y, copy=False)) + y_pred = convert_var_for_saving(np.array(y_pred, copy=False)) + + if out_fn is None: + out_fn = fn_template.format(batch_index=batch_index) + out_fn = os.path.join(out_dir, out_fn) + + with netCDF4.Dataset(out_fn, 'w') as ds: + dim_sample = ds.createDimension("dim_sample", y.shape[0]) + dim_channel = ds.createDimension("dim_channel", y.shape[1]) + dim_time_past = ds.createDimension("dim_time_past", x.shape[2]) + dim_time_future = ds.createDimension("dim_time_future", y.shape[2]) + dim_h = ds.createDimension("dim_h", y.shape[3]) + dim_w = ds.createDimension("dim_w", y.shape[4]) + dim_member = ds.createDimension("dim_member", y_pred.shape[5]) + var_params = {"zlib": True, "complevel": 1} + + var_fc = ds.createVariable( + "forecasts", y_pred.dtype, + ( + "dim_sample", "dim_channel", + "dim_time_future", "dim_h", "dim_w", "dim_member" + ), + **var_params + ) + var_fc[:] = y_pred + + var_obs_past = ds.createVariable( + "past_observations", x.dtype, + ("dim_sample", "dim_channel", "dim_time_past", "dim_h", "dim_w"), + **var_params + ) + var_obs_past[:] = x + + var_obs_future = ds.createVariable( + "future_observations", y.dtype, + ("dim_sample", "dim_channel", "dim_time_future", "dim_h", "dim_w"), + **var_params + ) + var_obs_future[:] = y + + +def load_batch(fn, decode=True, preproc_fc=None, **kwargs): + with netCDF4.Dataset(fn, 'r') as ds: + y_pred = np.array(ds["forecasts"][:], copy=False) + x = np.array(ds["past_observations"][:], copy=False) + y = np.array(ds["future_observations"][:], copy=False) + + if decode: + x = decode_saved_var_to_rainrate(x, **kwargs) + y = decode_saved_var_to_rainrate(y, **kwargs) + y_pred = decode_saved_var_to_rainrate( + y_pred, preproc=preproc_fc, **kwargs + ) + + return (x, y, y_pred) + + +def load_all_observations( + ensemble_dir, decode=True, preproc_fc=None, + timeframe='future', **kwargs +): + files = os.listdir(ensemble_dir) + obs = [] + for fn in sorted(files): + with netCDF4.Dataset(os.path.join(ensemble_dir, fn), 'r') as ds: + var = f"{timeframe}_observations" + x = np.array(ds[var][:], copy=False) + x = decode_saved_var_to_rainrate(x, **kwargs) + obs.append(x) + + obs = np.concatenate(obs, axis=0) + return obs diff --git a/ldcast/features/patches.py b/ldcast/features/patches.py new file mode 100644 index 0000000000000000000000000000000000000000..bb0aac8e9a3704c8a0a63b7f7fc0bb43843230e4 --- /dev/null +++ b/ldcast/features/patches.py @@ -0,0 +1,429 @@ +from datetime import datetime, timedelta +import os + +import dask +import netCDF4 +import numpy as np + +from .utils import average_pool + + +def patch_locations( + time_range, + patch_box, + patch_shape=(32,32), + interval=timedelta(minutes=5), + epoch=(1970,1,1) +): + patches = {} + t = time_range[0] + while t < time_range[1]: + patches[t] = [] + for pi in range(patch_box[0][0], patch_box[0][1]): + for pj in range(patch_box[1][0], patch_box[1][1]): + patches[t].append((pi,pj)) + patches[t] = np.array(patches[t]) + t += interval + + return patches + + +def save_patches_radar( + patches, archive_path, out_dir, + variables=("RZC", "CPCH"), + suffix="2020", + **kwargs +): + from ..datasets import mchradar + + source_vars = {} + mchradar_reader = mchradar.MCHRadarReader( + archive_path=archive_path, + variables=variables, + phys_values=False + ) + + ezc_nonzero_count_func = lambda x: np.count_nonzero((x >= 1) & (x<251)) + nonzero_count_func = { + "RZC": lambda x: np.count_nonzero(x > 1), + "CPCH": lambda x: np.count_nonzero(x > 1) + } + zero_value = {v: 0 for v in variables} + zero_value["RZC"] = 1 + zero_value["CPCH"] = 1 + + save_patches_all( + mchradar_reader, patches, variables, + nonzero_count_func, zero_value, out_dir, suffix, + source_vars=source_vars, min_nonzeros_to_include=5, + **kwargs + ) + + +def save_patches_dwdradar( + patches, archive_path, out_dir, + variables=("RV",), + suffix="2022", + patch_shape=(32,32), + **kwargs +): + from ..datasets import dwdradar + + source_vars = {} + dwdradar_reader = dwdradar.DWDRadarReader( + archive_path=archive_path, + variables=variables + ) + + patches_flt = {} + for t in sorted(patches): + if (t.hour==0) and (t.minute==0): + print(t) + + try: + data = dwdradar_reader.variable_for_time(t, "RV") + except FileNotFoundError: + continue + + patch_locs_time = [] + for (pi,pj) in patches[t]: + i0 = pi * patch_shape[0] + i1 = i0 + patch_shape[0] + j0 = pj * patch_shape[1] + j1 = j0 + patch_shape[1] + patch = data[i0:i1,j0:j1] + if np.isfinite(patch).all(): + patch_locs_time.append((pi,pj)) + + if patch_locs_time: + patches_flt[t] = np.array(patch_locs_time) + + print(len(patches),len(patches_flt)) + patches = patches_flt + + nonzero_count_func = {"RV": np.count_nonzero} + zero_value = {v: 0 for v in variables} + + save_patches_all( + dwdradar_reader, patches, variables, + nonzero_count_func, zero_value, out_dir, suffix, + source_vars=source_vars, min_nonzeros_to_include=5, + **kwargs + ) + + +def save_patches_ifs( + patches, archive_path, out_dir, + variables=( + #"rate-tp", "rate-cp", "t2m", "cape", "cin", + #"tclw", "tcwv", #, "rate-tpe" + #"u", "v" + "cin", + ), + suffix="2020", + lags=(0,12) +): + from ..datasets import ifsnwp + from .. import projection + + proj = projection.GridProjection(projection.ccs4_swiss_grid_area) + ifsnwp_reader = ifsnwp.IFSNWPReader( + proj, + archive_path=archive_path, + variables=variables, + lags=lags, + ) + + # we only get data for every hour, so modify patches + ifs_patches = { + dt: pset for (dt, pset) in patches.items() + if (dt.minute == dt.second == dt.microsecond == 0) + } + + variables_with_lag = [] + for lag in ifsnwp_reader.lags: + variables_with_lag.extend(f"{v}-{lag}" for v in variables) + + count_positive = lambda x: np.count_nonzero(x > 0) + all_nonzero = lambda x: np.prod(x.shape) + nonzero_count_func = { + "rate-tp": count_positive, + "rate-cp": count_positive, + "t2m": all_nonzero, + "cape": count_positive, + "cin": count_positive, + "tclw": count_positive, + "tcwv": count_positive, + "u": all_nonzero, + "v": all_nonzero, + "rate-tpe": count_positive, + } + nonzero_count_func = { + v: nonzero_count_func[v.rsplit("-", 1)[0]] + for v in variables_with_lag + } + postproc = { + f"cin-{lag}": lambda x: np.nan_to_num(x, nan=0.0, copy=False) + for lag in lags + } + zero_value = {v: 0 for v in variables_with_lag} + avg_pool = lambda x: average_pool(x, factor=8, missing=np.nan) + pool = {v: avg_pool for v in variables_with_lag} + + save_patches_all(ifsnwp_reader, ifs_patches, variables_with_lag, + nonzero_count_func, zero_value, out_dir, suffix, pool=pool, + postproc=postproc) + + +def save_patches_cosmo(patches, archive_path, out_dir, suffix="2020"): + from ..datasets import cosmonwp + + cosmonwp_reader = cosmonwp.COSMOCCS4Reader( + archive_path=archive_path, cache_size=6000) + + # we only get data for every hour, so modify patches + cosmo_patches = {} + for (dt,pset) in patches.items(): + dt0 = datetime(dt.year, dt.month, dt.day, dt.hour) + dt1 = dt0 + timedelta(hours=1) + if dt0 not in cosmo_patches: + cosmo_patches[dt0] = set() + if dt1 not in cosmo_patches: + cosmo_patches[dt1] = set() + cosmo_patches[dt0].update(pset) + cosmo_patches[dt1].update(pset) + + variables = [ + "CAPE_MU", "CIN_MU", "SLI", + "HZEROCL", "LCL_ML", "MCONV", "OMEGA", + "T_2M", "T_SO", "SOILTYP" + ] + count_positive = lambda x: np.count_nonzero(x>0) + all_nonzero = lambda x: np.prod(x.shape) + nonzero_count_func = { + "CAPE_MU": count_positive, + "CIN_MU": count_positive, + "SLI": all_nonzero, + "HZEROCL": count_positive, + "LCL_ML": count_positive, + "MCONV": all_nonzero, + "OMEGA": all_nonzero, + "T_2M": all_nonzero, + "T_SO": all_nonzero, + "SOILTYP": lambda x: np.count_nonzero(x!=5) + } + zero_value = {v: 0 for v in variables} + zero_value["SOILTYP"] = 5 + + save_patches_all(cosmonwp_reader, cosmo_patches, variables, + nonzero_count_func, zero_value, out_dir, suffix, pool=pool) + + +def save_patches_all( + reader, patches, variables, nonzero_count_func, zero_value, + out_dir, suffix, epoch=datetime(1970,1,1), postproc={}, scale=None, + pool={}, source_vars={}, parallel=False, min_nonzeros_to_include=1 +): + + def save_var(var_name): + src_name = source_vars.get(var_name, var_name) + + (patch_data, patch_coords, patch_times, + zero_patch_coords, zero_patch_times) = get_patches( + reader, src_name, patches, + nonzero_count_func=nonzero_count_func[var_name], + postproc=postproc.get(var_name), + pool=pool.get(var_name) + ) + try: + time = epoch + timedelta(seconds=int(patch_times[0])) + var_scale = reader.get_scale(time, var_name) + except (AttributeError, KeyError): + var_scale = None if (scale is None) else scale[var_name] + pass + + var_name = var_name.replace("_", "-") + out_fn = f"patches_{var_name}_{suffix}.nc" + out_path = os.path.join(out_dir, var_name) + os.makedirs(out_path, exist_ok=True) + out_fn = os.path.join(out_path, out_fn) + + save_patches( + patch_data, patch_coords, patch_times, + zero_patch_coords, zero_patch_times, out_fn, + zero_value=zero_value[var_name], scale=var_scale + ) + + if parallel: + save_var = dask.delayed(save_var) + + jobs = [save_var(v) for v in variables] + if parallel: + dask.compute(jobs, scheduler='threads') + + +def get_patches( + reader, variable, patches, + patch_shape=(32,32), nonzero_count_func=None, + epoch=datetime(1970,1,1), postproc=None, + pool=None, min_nonzeros_to_include=1 +): + num_patches = sum(len(patches[t]) for t in patches) + patch_data = [] + patch_coords = [] + patch_times = [] + zero_patch_coords = [] + zero_patch_times = [] + + if hasattr(reader, "phys_values"): + phys_values = reader.phys_values + + k = 0 + try: + if hasattr(reader, "phys_values"): + reader.phys_values = False + for (t, p_coord) in patches.items(): + try: + data = reader.variable_for_time(t, variable) + except (ValueError, FileNotFoundError, KeyError, OSError): + continue + + if postproc is not None: + data = postproc(data) + + time_sec = np.int64((t-epoch).total_seconds()) + for (pi, pj) in p_coord: + if k % 100000 == 0: + print("{}: {}/{}".format(t, k, num_patches)) + patch_box = data[ + pi*patch_shape[0]:(pi+1)*patch_shape[0], + pj*patch_shape[1]:(pj+1)*patch_shape[1], + ].copy() + is_nonzero = (nonzero_count_func is not None) and \ + (nonzero_count_func(patch_box) < min_nonzeros_to_include) + if is_nonzero: + zero_patch_coords.append((pi,pj)) + zero_patch_times.append(time_sec) + else: + if pool is not None: + patch_box = pool(patch_box) + patch_data.append(patch_box) + patch_coords.append((pi,pj)) + patch_times.append(time_sec) + k += 1 + + finally: + if hasattr(reader, "phys_values"): + reader.phys_values = phys_values + + if zero_patch_coords: + zero_patch_coords = np.stack(zero_patch_coords, axis=0).astype(np.uint16) + zero_patch_times = np.stack(zero_patch_times, axis=0) + else: + zero_patch_coords = np.zeros((0,2), dtype=np.uint16) + zero_patch_times = np.zeros((0,), dtype=np.int64) + patch_data = np.stack(patch_data, axis=0) + patch_coords = np.stack(patch_coords, axis=0).astype(np.uint16) + patch_times = np.stack(patch_times, axis=0) + + return (patch_data, patch_coords, patch_times, + zero_patch_coords, zero_patch_times) + + +def save_patches(patch_data, patch_coords, patch_times, + zero_patch_coords, zero_patch_times, out_fn, zero_value=0, scale=None): + + with netCDF4.Dataset(out_fn, 'w') as ds: + dim_patch = ds.createDimension("dim_patch", patch_data.shape[0]) + dim_zero_patch = ds.createDimension("dim_zero_patch", zero_patch_coords.shape[0]) + dim_coord = ds.createDimension("dim_coord", 2) + dim_height = ds.createDimension("dim_height", patch_data.shape[1]) + dim_width = ds.createDimension("dim_width", patch_data.shape[2]) + + var_args = {"zlib": True, "complevel": 1} + + chunksizes = (min(2**10, patch_data.shape[0]), patch_data.shape[1], patch_data.shape[2]) + var_patch = ds.createVariable("patches", patch_data.dtype, + ("dim_patch","dim_height","dim_width"), chunksizes=chunksizes, **var_args) + var_patch[:] = patch_data + + var_patch_coord = ds.createVariable("patch_coords", patch_coords.dtype, + ("dim_patch","dim_coord"), **var_args) + var_patch_coord[:] = patch_coords + + var_patch_time = ds.createVariable("patch_times", patch_times.dtype, + ("dim_patch",), **var_args) + var_patch_time[:] = patch_times + + var_zero_patch_coord = ds.createVariable("zero_patch_coords", zero_patch_coords.dtype, + ("dim_zero_patch","dim_coord"), **var_args) + var_zero_patch_coord[:] = zero_patch_coords + + var_zero_patch_time = ds.createVariable("zero_patch_times", zero_patch_times.dtype, + ("dim_zero_patch",), **var_args) + var_zero_patch_time[:] = zero_patch_times + + ds.zero_value = zero_value + + if scale is not None: + dim_scale = ds.createDimension("dim_scale", len(scale)) + var_scale = ds.createVariable("scale", scale.dtype, ("dim_scale",), **var_args) + var_scale[:] = scale + + +def load_patches(fn, in_memory=True): + if in_memory: + with open(fn, 'rb') as f: + ds_raw = f.read() + fn = None + else: + ds_raw = None + + with netCDF4.Dataset(fn, 'r', memory=ds_raw) as ds: + patch_data = { + "patches": np.array(ds["patches"]), + "patch_coords": np.array(ds["patch_coords"]), + "patch_times": np.array(ds["patch_times"]), + "zero_patch_coords": np.array(ds["zero_patch_coords"]), + "zero_patch_times": np.array(ds["zero_patch_times"]), + "zero_value": ds.zero_value + } + if "scale" in ds.variables: + patch_data["scale"] = np.array(ds["scale"]) + + return patch_data + + +def load_all_patches(patch_dir, var): + files = os.listdir(patch_dir) + jobs = [] + for fn in files: + file_var = fn.split("_")[1] + if file_var == var: + fn = os.path.join(patch_dir, fn) + jobs.append(dask.delayed(load_patches)(fn)) + + file_data = dask.compute(jobs, scheduler="processes")[0] + patch_data = {} + keys = ["patches", "patch_coords", "patch_times", + "zero_patch_coords", "zero_patch_times"] + for k in keys: + patch_data[k] = np.concatenate( + [fd[k] for fd in file_data], + axis=0 + ) + patch_data["zero_value"] = file_data[0]["zero_value"] + if "scale" in file_data[0]: + patch_data["scale"] = file_data[0]["scale"] + + return patch_data + + +def unpack_patches(patch_data): + return ( + patch_data["patches"], + patch_data["patch_coords"], + patch_data["patch_times"], + patch_data["zero_patch_coords"], + patch_data["zero_patch_times"] + ) diff --git a/ldcast/features/patches.py.save b/ldcast/features/patches.py.save new file mode 100644 index 0000000000000000000000000000000000000000..222da4796fa7fb282a1d02a62902ad84ef4bb22c --- /dev/null +++ b/ldcast/features/patches.py.save @@ -0,0 +1,431 @@ +from datetime import datetime, timedelta +import os + +import dask +import netCDF4 +import numpy as np + +from .utils import average_pool + + +def patch_locations( + time_range, + patch_box, + patch_shape=(32,32), + interval=timedelta(minutes=5), + epoch=(1970,1,1) +): + patches = {} + t = time_range[0] + while t < time_range[1]: + patches[t] = [] + for pi in range(patch_box[0][0], patch_box[0][1]): + for pj in range(patch_box[1][0], patch_box[1][1]): + patches[t].append((pi,pj)) + patches[t] = np.array(patches[t]) + t += interval + + return patches + + +def save_patches_radar( + patches, archive_path, out_dir, + variables=("RZC", "CPCH"), + suffix="2020", + **kwargs +): + from ..datasets import mchradar + + source_vars = {} + mchradar_reader = mchradar.MCHRadarReader( + archive_path=archive_path, + variables=variables, + phys_values=False + ) + + ezc_nonzero_count_func = lambda x: np.count_nonzero((x >= 1) & (x<251)) + nonzero_count_func = { + "RZC": lambda x: np.count_nonzero(x > 1), + "CPCH": lambda x: np.count_nonzero(x > 1) + } + zero_value = {v: 0 for v in variables} + zero_value["RZC"] = 1 + zero_value["CPCH"] = 1 + + save_patches_all( + mchradar_reader, patches, variables, + nonzero_count_func, zero_value, out_dir, suffix, + source_vars=source_vars, min_nonzeros_to_include=5, + **kwargs + ) + + +def save_patches_dwdradar( + patches, archive_path, out_dir, + variables=("RV",), + suffix="2022", + patch_shape=(32,32), + **kwargs +): + from ..datasets import dwdradar + + source_vars = {} + dwdradar_reader = dwdradar.DWDRadarReader( + archive_path=archive_path, + variables=variables + ) + + patches_flt = {} + for t in sorted(patches): + if (t.hour==0) and (t.minute==0): + print(t) + + try: + data = dwdradar_reader.variable_for_time(t, "RV") + except FileNotFoundError: + continue + + patch_locs_time = [] + for (pi,pj) in patches[t]: + i0 = pi * patch_shape[0] + i1 = i0 + patch_shape[0] + j0 = pj * patch_shape[1] + j1 = j0 + patch_shape[1] + patch = data[i0:i1,j0:j1] + if np.isfinite(patch).all(): + patch_locs_time.append((pi,pj)) + + if patch_locs_time: + patches_flt[t] = np.array(patch_locs_time) + + print(len(patches),len(patches_flt)) + patches = patches_flt + + nonzero_count_func = {"RV": np.count_nonzero} + zero_value = {v: 0 for v in variables} + + save_patches_all( + dwdradar_reader, patches, variables, + nonzero_count_func, zero_value, out_dir, suffix, + source_vars=source_vars, min_nonzeros_to_include=5, + **kwargs + ) + + +def save_patches_ifs( + patches, archive_path, out_dir, + variables=( + #"rate-tp", "rate-cp", "t2m", "cape", "cin", + #"tclw", "tcwv", #, "rate-tpe" + #"u", "v" + "cin", + ), + suffix="2020", + lags=(0,12) +): + from ..datasets import ifsnwp + from .. import projection + + proj = projection.GridProjection(projection.ccs4_swiss_grid_area) + ifsnwp_reader = ifsnwp.IFSNWPReader( + proj, + archive_path=archive_path, + variables=variables, + lags=lags, + ) + + # we only get data for every hour, so modify patches + ifs_patches = { + dt: pset for (dt, pset) in patches.items() + if (dt.minute == dt.second == dt.microsecond == 0) + } + + variables_with_lag = [] + for lag in ifsnwp_reader.lags: + variables_with_lag.extend(f"{v}-{lag}" for v in variables) + + count_positive = lambda x: np.count_nonzero(x > 0) + all_nonzero = lambda x: np.prod(x.shape) + nonzero_count_func = { + "rate-tp": count_positive, + "rate-cp": count_positive, + "t2m": all_nonzero, + "cape": count_positive, + "cin": count_positive, + "tclw": count_positive, + "tcwv": count_positive, + "u": all_nonzero, + "v": all_nonzero, + "rate-tpe": count_positive, + } + nonzero_count_func = { + v: nonzero_count_func[v.rsplit("-", 1)[0]] + for v in variables_with_lag + } + postproc = { + f"cin-{lag}": lambda x: np.nan_to_num(x, nan=0.0, copy=False) + for lag in lags + } + zero_value = {v: 0 for v in variables_with_lag} + avg_pool = lambda x: average_pool(x, factor=8, missing=np.nan) + pool = {v: avg_pool for v in variables_with_lag} + + save_patches_all(ifsnwp_reader, ifs_patches, variables_with_lag, + nonzero_count_func, zero_value, out_dir, suffix, pool=pool, + postproc=postproc) + + +def save_patches_cosmo(patches, archive_path, out_dir, suffix="2020"): + from ..datasets import cosmonwp + + cosmonwp_reader = cosmonwp.COSMOCCS4Reader( + archive_path=archive_path, cache_size=6000) + + # we only get data for every hour, so modify patches + cosmo_patches = {} + for (dt,pset) in patches.items(): + dt0 = datetime(dt.year, dt.month, dt.day, dt.hour) + dt1 = dt0 + timedelta(hours=1) + if dt0 not in cosmo_patches: + cosmo_patches[dt0] = set() + if dt1 not in cosmo_patches: + cosmo_patches[dt1] = set() + cosmo_patches[dt0].update(pset) + cosmo_patches[dt1].update(pset) + + variables = [ + "CAPE_MU", "CIN_MU", "SLI", + "HZEROCL", "LCL_ML", "MCONV", "OMEGA", + "T_2M", "T_SO", "SOILTYP" + ] + count_positive = lambda x: np.count_nonzero(x>0) + all_nonzero = lambda x: np.prod(x.shape) + nonzero_count_func = { + "CAPE_MU": count_positive, + "CIN_MU": count_positive, + "SLI": all_nonzero, + "HZEROCL": count_positive, + "LCL_ML": count_positive, + "MCONV": all_nonzero, + "OMEGA": all_nonzero, + "T_2M": all_nonzero, + "T_SO": all_nonzero, + "SOILTYP": lambda x: np.count_nonzero(x!=5) + } + zero_value = {v: 0 for v in variables} + zero_value["SOILTYP"] = 5 + + save_patches_all(cosmonwp_reader, cosmo_patches, variables, + nonzero_count_func, zero_value, out_dir, suffix, pool=pool) + + +def save_patches_all( + reader, patches, variables, nonzero_count_func, zero_value, + out_dir, suffix, epoch=datetime(1970,1,1), postproc={}, scale=None, + pool={}, source_vars={}, parallel=False, min_nonzeros_to_include=1 +): + + def save_var(var_name): + src_name = source_vars.get(var_name, var_name) + + (patch_data, patch_coords, patch_times, + zero_patch_coords, zero_patch_times) = get_patches( + reader, src_name, patches, + nonzero_count_func=nonzero_count_func[var_name], + postproc=postproc.get(var_name), + pool=pool.get(var_name) + ) + try: + time = epoch + timedelta(seconds=int(patch_times[0])) + var_scale = reader.get_scale(time, var_name) + except (AttributeError, KeyError): + var_scale = None if (scale is None) else scale[var_name] + pass + + var_name = var_name.replace("_", "-") + out_fn = f"patches_{var_name}_{suffix}.nc" + out_path = os.path.join(out_dir, var_name) + os.makedirs(out_path, exist_ok=True) + out_fn = os.path.join(out_path, out_fn) + + save_patches( + patch_data, patch_coords, patch_times, + zero_patch_coords, zero_patch_times, out_fn, + zero_value=zero_value[var_name], scale=var_scale + ) + + if parallel: + save_var = dask.delayed(save_var) + + jobs = [save_var(v) for v in variables] + if parallel: + dask.compute(jobs, scheduler='threads') + + +def get_patches( + reader, variable, patches, + patch_shape=(32,32), nonzero_count_func=None, + epoch=datetime(1970,1,1), postproc=None, + pool=None, min_nonzeros_to_include=1 +): + num_patches = sum(len(patches[t]) for t in patches) + patch_data = [] + patch_coords = [] + patch_times = [] + zero_patch_coords = [] + zero_patch_times = [] + + if hasattr(reader, "phys_values"): + phys_values = reader.phys_values + + k = 0 + try: + if hasattr(reader, "phys_values"): + reader.phys_values = False + for (t, p_coord) in patches.items(): + try: + data = reader.variable_for_time(t, variable) + except (ValueError, FileNotFoundError, KeyError, OSError): + continue + + if postproc is not None: + data = postproc(data) + + time_sec = np.int64((t-epoch).total_seconds()) + for (pi, pj) in p_coord: + if k % 100000 == 0: + print("{}: {}/{}".format(t, k, num_patches)) + patch_box = data[ + pi*patch_shape[0]:(pi+1)*patch_shape[0], + pj*patch_shape[1]:(pj+1)*patch_shape[1], + ].copy() + is_nonzero = (nonzero_count_func is not None) and \ + (nonzero_count_func(patch_box) < min_nonzeros_to_include) + if is_nonzero: + zero_patch_coords.append((pi,pj)) + zero_patch_times.append(time_sec) + else: + if pool is not None: + patch_box = pool(patch_box) + patch_data.append(patch_box) + patch_coords.append((pi,pj)) + patch_times.append(time_sec) + k += 1 + + finally: + if hasattr(reader, "phys_values"): + reader.phys_values = phys_values + + if zero_patch_coords: + zero_patch_coords = np.stack(zero_patch_coords, axis=0).astype(np.uint16) + zero_patch_times = np.stack(zero_patch_times, axis=0) + else: + zero_patch_coords = np.zeros((0,2), dtype=np.uint16) + zero_patch_times = np.zeros((0,), dtype=np.int64) + patch_data = np.stack(patch_data, axis=0) + patch_coords = np.stack(patch_coords, axis=0).astype(np.uint16) + patch_times = np.stack(patch_times, axis=0) + + return (patch_data, patch_coords, patch_times, + zero_patch_coords, zero_patch_times) + + +def save_patches(patch_data, patch_coords, patch_times, + zero_patch_coords, zero_patch_times, out_fn, zero_value=0, scale=None): + + with netCDF4.Dataset(out_fn, 'w') as ds: + dim_patch = ds.createDimension("dim_patch", patch_data.shape[0]) + dim_zero_patch = ds.createDimension("dim_zero_patch", zero_patch_coords.shape[0]) + dim_coord = ds.createDimension("dim_coord", 2) + dim_height = ds.createDimension("dim_height", patch_data.shape[1]) + dim_width = ds.createDimension("dim_width", patch_data.shape[2]) + + var_args = {"zlib": True, "complevel": 1} + + chunksizes = (min(2**10, patch_data.shape[0]), patch_data.shape[1], patch_data.shape[2]) + var_patch = ds.createVariable("patches", patch_data.dtype, + ("dim_patch","dim_height","dim_width"), chunksizes=chunksizes, **var_args) + var_patch[:] = patch_data + + var_patch_coord = ds.createVariable("patch_coords", patch_coords.dtype, + ("dim_patch","dim_coord"), **var_args) + var_patch_coord[:] = patch_coords + + var_patch_time = ds.createVariable("patch_times", patch_times.dtype, + ("dim_patch",), **var_args) + var_patch_time[:] = patch_times + + var_zero_patch_coord = ds.createVariable("zero_patch_coords", zero_patch_coords.dtype, + ("dim_zero_patch","dim_coord"), **var_args) + var_zero_patch_coord[:] = zero_patch_coords + + var_zero_patch_time = ds.createVariable("zero_patch_times", zero_patch_times.dtype, + ("dim_zero_patch",), **var_args) + var_zero_patch_time[:] = zero_patch_times + + ds.zero_value = zero_value + + if scale is not None: + dim_scale = ds.createDimension("dim_scale", len(scale)) + var_scale = ds.createVariable("scale", scale.dtype, ("dim_scale",), **var_args) + var_scale[:] = scale + + +def load_patches(fn, in_memory=True): + if in_memory: + with open(fn, 'rb') as f: + ds_raw = f.read() + fn = None + else: + ds_raw = None + + with netCDF4.Dataset(fn, 'r', memory=ds_raw) as ds: + patch_data = { + #"patches": np.array(ds["patches"]), + #"patch_coords": np.array(ds["patch_coords"]), + #"patch_times": np.array(ds["patch_times"]), + #"zero_patch_coords": np.array(ds["zero_patch_coords"]), + #"zero_patch_times": np.array(ds["zero_patch_times"]), + "zero_value": 1, + "pr": np.array(ds["pr"]), + } + if "scale" in ds.variables: + patch_data["scale"] = np.array(ds["scale"]) + + return patch_data + + +def load_all_patches(patch_dir, var): + files = os.listdir(patch_dir) + jobs = [] + for fn in files: + file_var = fn.split("_")[1] + if file_var == var: + fn = os.path.join(patch_dir, fn) + jobs.append(dask.delayed(load_patches)(fn)) + + file_data = dask.compute(jobs, scheduler="processes")[0] + patch_data = {} + #keys = ["patches", "patch_coords", "patch_times", + # "zero_patch_coords", "zero_patch_times"] + keys = ["pr"] + for k in keys: + patch_data[k] = np.concatenate( + [fd[k] for fd in file_data], + axis=0 + ) + patch_data["zero_value"] = file_data[0]["zero_value"] + if "scale" in file_data[0]: + patch_data["scale"] = file_data[0]["scale"] + + return patch_data + + +def unpack_patches(patch_data): + return ( + patch_data["patches"], + patch_data["patch_coords"], + patch_data["patch_times"], + patch_data["zero_patch_coords"], + patch_data["zero_patch_times"] + ) diff --git a/ldcast/features/sampling.py b/ldcast/features/sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..376335cae7552f431c90edcd6766a669e25da684 --- /dev/null +++ b/ldcast/features/sampling.py @@ -0,0 +1,215 @@ +from bisect import bisect_left +import multiprocessing + +import dask +from numba import njit, prange, types +from numba.typed import Dict +import numpy as np + +from .patches import unpack_patches + + +class EqualFrequencySampler: + def __init__( + self, bins, patch_data, patch_index, + sample_shape, time_range_valid, time_range_sampling=None, + timestep_secs=5*60, + random_seed=None, preselected_samples=None + ): + binned_patches = bin_classify_patches_parallel( + bins, + *unpack_patches(patch_data), + zero_value=patch_data.get("zero_value", 0), + scale=patch_data.get("scale") + ) + complete_ind = indices_with_complete_sample( + patch_index, sample_shape, time_range_valid, timestep_secs + ) + if time_range_sampling is None: + time_range_sampling = time_range_valid + self.starting_ind = [ + starting_indices_for_centers( + p, complete_ind, sample_shape, time_range_sampling, timestep_secs + ) + for p in binned_patches + ] + self.num_bins = len(self.starting_ind) + self.rng = np.random.RandomState(seed=random_seed) + self.preselected_samples = preselected_samples + self.current_ind = np.array([len(ind) for ind in self.starting_ind]) + + def get_bin_sample(self, bin_ind): + patches = self.starting_ind[bin_ind] + sample_ind = self.current_ind[bin_ind] + if sample_ind >= patches.shape[0]: + self.rng.shuffle(patches) + sample_ind = self.current_ind[bin_ind] = 0 + else: + self.current_ind[bin_ind] += 1 + return patches[sample_ind,:] + + def __call__(self, num): + # sample each bin with equal probability + bins = self.rng.randint(self.num_bins, size=num) + coords = np.stack( + [self.get_bin_sample(b) for b in bins], + axis=0 + ) + return coords + + +def bin_classify_patches( + bins, patches, patch_coords, patch_times, + zero_patch_coords, zero_patch_times, + zero_value=0, metric_func=None, + scale=None, +): + if metric_func is None: + def metric_func(x): + xm = np.percentile(x, 99, axis=(1,2)) + if np.issubdtype(x.dtype, np.integer): + xm = xm.round() + return xm.astype(x.dtype) + + binned_patches = [[] for _ in range(len(bins)+1)] + + def find_bin(value): + return bisect_left(bins, value) + + zero_bin = find_bin(zero_value if scale is None else scale[zero_value]) + for (t,(pi,pj)) in zip(zero_patch_times, zero_patch_coords): + binned_patches[zero_bin].append((t,pi,pj)) + + patch_metrics = metric_func(patches) + if scale is not None: + patch_metrics = scale[patch_metrics] + for (metric,t,(pi,pj)) in zip(patch_metrics, patch_times, patch_coords): + patch_bin = find_bin(metric) + binned_patches[patch_bin].append((t,pi,pj)) + + for i in range(len(binned_patches)): + if binned_patches[i]: + binned_patches[i] = np.array(binned_patches[i]) + else: + binned_patches[i] = np.zeros((0,3), dtype=np.int64) + + return binned_patches + + +def bin_classify_patches_parallel( + bins, patches, patch_coords, patch_times, + zero_patch_coords, zero_patch_times, + zero_value=0, metric_func=None, + scale=None, +): + num_patches = patches.shape[0] + num_zeros = zero_patch_coords.shape[0] + num_procs = multiprocessing.cpu_count() + + tasks = [] + for p in range(num_procs): + pk0 = int(round(num_patches*p/num_procs)) + pk1 = int(round(num_patches*(p+1)/num_procs)) + zk0 = int(round(num_zeros*p/num_procs)) + zk1 = int(round(num_zeros*(p+1)/num_procs)) + + task = dask.delayed(bin_classify_patches)( + bins, + patches[pk0:pk1,...], patch_coords[pk0:pk1,...], + patch_times[pk0:pk1], + zero_patch_coords[zk0:zk1,...], zero_patch_times[zk0:zk1], + zero_value=zero_value, metric_func=metric_func, + scale=scale + ) + tasks.append(task) + + chunked_bins = dask.compute(tasks, scheduler="threads")[0] + + n_bins = len(chunked_bins[0]) + binned_patches = [ + np.concatenate([cb[i] for cb in chunked_bins], axis=0) + for i in range(n_bins) + ] + return binned_patches + + +def indices_with_complete_sample( + patch_index, sample_shape, time_range, timestep_secs +): + """Check which locations will give a sample without missing data. + """ + ind = np.array(list(patch_index.patch_index.keys())) + t0 = ind[:,0] + i0 = ind[:,1] + j0 = ind[:,2] + n = ind.shape[0] + complete = np.ones(n, dtype=bool) + # we use this dict like a set - numba doesn't support typed sets + complete_ind = Dict.empty( + key_type=types.UniTuple(types.int64, 3), + value_type=types.uint8 + ) + + @njit(parallel=True) # many nested loops, numba optimization needed + def check_complete(index, complete, complete_ind): + for k in prange(n): + for ts in range(*time_range): + t = t0[k] + ts*timestep_secs + for di in range(sample_shape[0]): + i = i0[k] + di + for dj in range(sample_shape[1]): + j = j0[k] + dj + if (t,i,j) not in index: + complete[k] = False + + for k in range(n): # no prange: can't set dict items in parallel + if complete[k]: + complete_ind[(t0[k],i0[k],j0[k])] = np.uint8(0) + + check_complete(patch_index.patch_index, complete, complete_ind) + + return complete_ind + + +def starting_indices_for_centers( + centers, complete_ind, sample_shape, time_range, timestep_secs +): + """Determine a complete list of sample indices that + contain one or more of the centerpoints. + """ + + @njit + def find_indices(centers, starting_ind, complete_ind): + for k in range(centers.shape[0]): + t0 = centers[k,0] + i0 = centers[k,1] + j0 = centers[k,2] + for ts in range(*time_range): + t = t0 - ts*timestep_secs # note minus signs in (t,i,j) + for di in range(sample_shape[0]): + i = i0 - di + for dj in range(sample_shape[1]): + j = j0 - dj + if (t,i,j) in complete_ind: + starting_ind[(t,i,j)] = np.uint8(0) + + num_chunks = multiprocessing.cpu_count() + + @dask.delayed + def chunk(i): + starting_ind = Dict.empty( + key_type=types.UniTuple(types.int64, 3), + value_type=types.uint8 + ) + k0 = int(round(centers.shape[0] * (i / num_chunks))) + k1 = int(round(centers.shape[0] * ((i+1) / num_chunks))) + find_indices(centers[k0:k1,...], starting_ind, complete_ind) + return starting_ind + + jobs = [chunk(i) for i in range(num_chunks)] + starting_ind = dask.compute(jobs, scheduler='threads')[0] + starting_ind = np.concatenate( + [np.array(list(st_ind.keys())) for st_ind in starting_ind if st_ind], + axis=0 + ) + return starting_ind diff --git a/ldcast/features/split.py b/ldcast/features/split.py new file mode 100644 index 0000000000000000000000000000000000000000..d176d31e803c1b480072cada46c3e1de66737f77 --- /dev/null +++ b/ldcast/features/split.py @@ -0,0 +1,165 @@ +from bisect import bisect_left + +import numpy as np +import pytorch_lightning as pl +from torch.utils.data import DataLoader + +from . import batch + + +def get_chunks( + primary_raw, valid_frac=0.1, test_frac=0.1, + chunk_seconds=2*24*60*60, random_seed=None +): + t0 = min( + primary_raw["patch_times"][0], + primary_raw["zero_patch_times"][0] + ) + t1 = max( + primary_raw["patch_times"][-1], + primary_raw["zero_patch_times"][-1] + )+1 + + rng = np.random.RandomState(seed=random_seed) + chunk_limits = np.arange(t0,t1,chunk_seconds) + num_chunks = len(chunk_limits)-1 + + chunk_ind = np.arange(num_chunks) + rng.shuffle(chunk_ind) + i_valid = int(round(num_chunks * valid_frac)) + i_test = i_valid + int(round(num_chunks * test_frac)) + chunk_ind = { + "valid": chunk_ind[:i_valid], + "test": chunk_ind[i_valid:i_test], + "train": chunk_ind[i_test:] + } + def get_chunk_limits(chunk_ind_split): + return sorted( + (chunk_limits[i], chunk_limits[i+1]) + for i in chunk_ind_split + ) + chunks = { + split: get_chunk_limits(chunk_ind_split) + for (split, chunk_ind_split) in chunk_ind.items() + } + return chunks + + +def train_valid_test_split( + raw_data, primary_raw_var, chunks=None, **kwargs +): + if chunks is None: + primary = raw_data[primary_raw_var] + chunks = get_chunks(primary, **kwargs) + + def split_chunks_from_array(x, chunks_split, times): + n = 0 + chunk_ind = [] + for (t0,t1) in chunks_split: + k0 = bisect_left(times, t0) + k1 = bisect_left(times, t1) + n += k1 - k0 + chunk_ind.append((k0,k1)) + + shape = (n,) + x.shape[1:] + x_chunk = np.empty_like(x, shape=shape) + + j0 = 0 + for (k0,k1) in chunk_ind: + j1 = j0 + (k1-k0) + x_chunk[j0:j1,...] = x[k0:k1,...] + j0 = j1 + + return x_chunk + + split_raw_data = { + split: {var: {} for var in raw_data} + for split in chunks + } + + for (var, raw_data_var) in raw_data.items(): + for (split, chunks_split) in chunks.items(): + + #split_raw_data[split][var]["patches"] = \ + # split_chunks_from_array( + # raw_data_var["patches"], chunks_split, + # raw_data_var["patch_times"] + # ) + #split_raw_data[split][var]["patch_coords"] = \ + # split_chunks_from_array( + # raw_data_var["patch_coords"], chunks_split, + # raw_data_var["patch_times"] + # ) + #split_raw_data[split][var]["patch_times"] = \ + # split_chunks_from_array( + # raw_data_var["patch_times"], chunks_split, + # raw_data_var["patch_times"] + # ) + #split_raw_data[split][var]["zero_patch_coords"] = \ + # split_chunks_from_array( + # raw_data_var["zero_patch_coords"], chunks_split, + # raw_data_var["zero_patch_times"] + # ) + #split_raw_data[split][var]["zero_patch_times"] = \ + # split_chunks_from_array( + # raw_data_var["zero_patch_times"], chunks_split, + # raw_data_var["zero_patch_times"] + # ) + + added_keys = set(split_raw_data[split][var].keys()) + missing_keys = set(raw_data[var].keys()) - added_keys + for k in missing_keys: + split_raw_data[split][var][k] = raw_data[var][k] + + return (split_raw_data, chunks) + + +class DataModule(pl.LightningDataModule): + def __init__( + self, + variables, raw, predictors, target, primary_var, + sampling_bins, sampler_file, + batch_size=8, + train_epoch_size=10, valid_epoch_size=2, test_epoch_size=10, + valid_seed=None, test_seed=None, + **kwargs + ): + super().__init__() + self.batch_gen = { + split: batch.BatchGenerator( + variables, raw_var, predictors, target, primary_var, + sampling_bins=sampling_bins, batch_size=batch_size, + sampler_file=sampler_file.get(split), + augment=(split=="train"), + **kwargs + ) + for (split,raw_var) in raw.items() + } + self.datasets = {} + if "train" in self.batch_gen: + self.datasets["train"] = batch.StreamBatchDataset( + self.batch_gen["train"], train_epoch_size + ) + if "valid" in self.batch_gen: + self.datasets["valid"] = batch.DeterministicBatchDataset( + self.batch_gen["valid"], valid_epoch_size, random_seed=valid_seed + ) + if "test" in self.batch_gen: + self.datasets["test"] = batch.DeterministicBatchDataset( + self.batch_gen["test"], test_epoch_size, random_seed=test_seed + ) + + def dataloader(self, split): + return DataLoader( + self.datasets[split], batch_size=None, + pin_memory=True, num_workers=0 + ) + + def train_dataloader(self): + return self.dataloader("train") + + def val_dataloader(self): + return self.dataloader("valid") + + def test_dataloader(self): + return self.dataloader("test") diff --git a/ldcast/features/transform.py b/ldcast/features/transform.py new file mode 100644 index 0000000000000000000000000000000000000000..3e5d3b7f1881f02ae9def77030d85ff7eebfffc0 --- /dev/null +++ b/ldcast/features/transform.py @@ -0,0 +1,296 @@ +import concurrent.futures +import multiprocessing + +from numba import njit, prange +import numpy as np +from scipy.ndimage import convolve + + +def quick_cast(x, y): + num_threads = multiprocessing.cpu_count() + with concurrent.futures.ThreadPoolExecutor(num_threads) as executor: + futures = {} + limits = np.linspace(0, x.shape[0], num_threads+1).round().astype(int) + def _cast(k0,k1): + y[k0:k1,...] = x[k0:k1,...] + for k in range(len(limits)-1): + args = (_cast, limits[k], limits[k+1]) + futures[executor.submit(*args)] = k + concurrent.futures.wait(futures) + + +def cast(dtype=np.float16): + xc = None + def transform(raw): + nonlocal xc + if (xc is None) or (xc.shape != raw.shape): + xc = np.empty_like(raw, dtype=dtype) + quick_cast(raw, xc) + return xc + return transform + + +@njit(parallel=True) +def scale_array(in_arr, out_arr, scale): + in_arr = in_arr.ravel() + out_arr = out_arr.ravel() + for i in prange(in_arr.shape[0]): + out_arr[i] = scale[in_arr[i]] + +# NumPy version +#def scale_array(in_arr, out_arr, scale): +# out_arr[:] = scale[in_arr] + +def normalize(mean=0.0, std=1.0, dtype=np.float32): + scaled = scaled_dt = None + + def transform(raw): + nonlocal scaled, scaled_dt + if (scaled is None) or (scaled.shape != raw.shape): + scaled = np.empty_like(raw, dtype=np.float32) + scaled_dt = np.empty_like(raw, dtype=dtype) + normalize_array(raw, scaled, mean, std) + + if dtype == np.float32: + return scaled + else: + quick_cast(scaled, scaled_dt) + return scaled_dt + + return transform + + +def normalize_threshold(mean=0.0, std=1.0, threshold=0.0, fill_value=0.0, log=False): + scaled = None + + def transform(raw): + nonlocal scaled + if (scaled is None) or (scaled.shape != raw.shape): + scaled = np.empty_like(raw, dtype=np.float32) + normalize_threshold_array(raw, scaled, mean, std, threshold, fill_value, log=log) + + return scaled + + return transform + + +def scale_log_norm(scale, threshold=None, missing_value=None, + fill_value=0, mean=0.0, std=1.0, dtype=np.float32): + + log_scale = np.log10(scale, where=scale>0).astype(np.float32) + if threshold is not None: + log_scale[log_scale < np.log10(threshold)] = np.log10(fill_value) + if missing_value is not None: + log_scale[missing_value] = np.log10(fill_value) + log_scale[~np.isfinite(log_scale)] = np.log10(fill_value) + log_scale -= mean + log_scale /= std + scaled = scaled_dt = None + + def transform(raw): + nonlocal scaled, scaled_dt + if (scaled is None) or (scaled.shape != raw.shape): + scaled = np.empty_like(raw, dtype=np.float32) + scaled_dt = np.empty_like(raw, dtype=dtype) + scale_array(raw, scaled, log_scale) + + if dtype == np.float32: + return scaled + else: + quick_cast(scaled, scaled_dt) + return scaled_dt + + return transform + + +def combine(transforms, memory_format="channels_first", dim=3): + #combined = None + channels_axis = 1 if (memory_format == "channels_first") else -1 + + def transform(*raw): + #nonlocal combined + transformed = [t(r) for (t, r) in zip(transforms, raw)] + for i in range(len(transformed)): + if transformed[i].ndim == dim + 1: + transformed[i] = np.expand_dims(transformed[i], channels_axis) + + return np.concatenate(transformed, axis=channels_axis) + + return transform + + +class Antialiasing: + def __init__(self): + (x,y) = np.mgrid[-2:3,-2:3] + self.kernel = np.exp(-0.5*(x**2+y**2)/(0.5**2)) + self.kernel /= self.kernel.sum() + self.edge_factors = {} + self.img_smooth = {} + num_threads = multiprocessing.cpu_count() + self.executor = concurrent.futures.ThreadPoolExecutor(num_threads) + + def __call__(self, img): + img_shape = img.shape[-2:] + if img_shape not in self.edge_factors: + s = convolve(np.ones(img_shape, dtype=np.float32), + self.kernel, mode="constant") + s = 1.0/s + self.edge_factors[img_shape] = s + else: + s = self.edge_factors[img_shape] + + if img.shape not in self.img_smooth: + img_smooth = np.empty_like(img) + self.img_smooth[img_shape] = img_smooth + else: + img_smooth = self.img_smooth[img_shape] + + def _convolve_frame(i,j): + convolve(img[i,j,:,:], self.kernel, + mode="constant", output=img_smooth[i,j,:,:]) + img_smooth[i,j,:,:] *= s + + futures = [] + for i in range(img.shape[0]): + for j in range(img.shape[1]): + args = (_convolve_frame, i, j) + futures.append(self.executor.submit(*args)) + concurrent.futures.wait(futures) + + return img_smooth + + +def default_rainrate_transform(scale): + scaling = scale_log_norm( + scale, threshold=0.1, fill_value=0.02, + mean=-0.051, std=0.528, dtype=np.float32 + ) + antialiasing = Antialiasing() + def transform(raw): + x = scaling(raw) + return antialiasing(x) + return transform + + +def scale_norm(scale, threshold=None, missing_value=None, + fill_value=0, mean=0.0, std=1.0, dtype=np.float32): + + scale = scale.astype(np.float32).copy() + scale[np.isnan(scale)] = fill_value + if threshold is not None: + scale[scale < threshold] = fill_value + if missing_value is not None: + missing_value = np.atleast_1d(missing_value) + for m in missing_value: + scale[m] = fill_value + scale -= mean + scale /= std + scaled = scaled_dt = None + + def transform(raw): + nonlocal scaled, scaled_dt + if (scaled is None) or (scaled.shape != raw.shape): + scaled = np.empty_like(raw, dtype=np.float32) + scaled_dt = np.empty_like(raw, dtype=dtype) + scale_array(raw, scaled, scale) + + if dtype == np.float32: + return scaled + else: + quick_cast(scaled, scaled_dt) + return scaled_dt + + return transform + + +@njit(parallel=True) +def threshold_array(in_arr, out_arr, threshold): + in_arr = in_arr.ravel() + out_arr = out_arr.ravel() + for i in prange(in_arr.shape[0]): + out_arr[i] = np.float32(in_arr[i] >= threshold) + + +def one_hot(values): + translation = np.zeros(max(values)+1, dtype=int) + num_categories = len(values) + for (i,v) in enumerate(values): + translation[v] = i + onehot = onehot_dt = None + + def transform(raw): + nonlocal onehot, onehot_dt + if (onehot is None) or (onehot.shape[:-1] != raw.shape): + onehot = np.empty(raw.shape+(num_categories,), + dtype=np.float32) + onehot = np.empty(raw.shape+(num_categories,), + dtype=np.uint8) + onehot_transform(raw, onehot, translation) + quick_cast(onehot, onehot_dt) + + return onehot + + return transform + + +@njit(parallel=True) +def onehot_transform(in_arr, out_arr, translation): + for k in prange(in_arr.shape[0]): + out_arr[k,...] = 0.0 + for t in range(in_arr.shape[1]): + for i in range(in_arr.shape[2]): + for j in range(in_arr.shape[3]): + ind = np.uint64(in_arr[k,t,i,j]) + c = translation[ind] + out_arr[k,t,i,j,c] = 1.0 + + +@njit(parallel=True) +def normalize_array(in_arr, out_arr, mean, std): + mean = np.float32(mean) + inv_std = np.float32(1.0/std) + in_arr = in_arr.ravel() + out_arr = out_arr.ravel() + for i in prange(in_arr.shape[0]): + out_arr[i] = (in_arr[i]-mean)*inv_std + + +@njit(parallel=True) +def normalize_threshold_array( + in_arr, out_arr, + mean, std, + threshold, fill_value, log=False +): + mean = np.float32(mean) + inv_std = np.float32(1.0/std) + threshold = np.float32(threshold) + fill_value = np.float32(fill_value) + in_arr = in_arr.ravel() + out_arr = out_arr.ravel() + for i in prange(in_arr.shape[0]): + x = in_arr[i] + if x < threshold: + x = fill_value + if log: + x = np.log10(x) + out_arr[i] = (x-mean)*inv_std + + +# NumPy version +#def threshold_array(in_arr, out_arr, threshold): +# out_arr[:] = (in_arr >= threshold).astype(np.float32) + + +def R_threshold(scale, threshold): + thresholded = None + scale_treshold = np.nanargmax(scale > threshold) + + def transform(rzc_raw): + nonlocal thresholded + if (thresholded is None) or (thresholded.shape != rzc_raw.shape): + thresholded = np.empty_like(rzc_raw, dtype=np.float32) + threshold_array(rzc_raw, thresholded, scale_treshold) + + return thresholded + + return transform diff --git a/ldcast/features/utils.py b/ldcast/features/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..78eceaf852eef4f5581fd682404df4e57050c425 --- /dev/null +++ b/ldcast/features/utils.py @@ -0,0 +1,136 @@ +from numba import njit, prange +import numpy as np +from scipy.signal import convolve + + +def log_scale_with_zero(range, n=65536, dtype=np.float32): + scale = np.linspace(np.log10(range[0]), np.log10(range[1]), n-1) + scale = np.hstack((0, 10**scale)).astype(dtype) + return scale + + +def log_quantize_with_zero(x, range, n=65536, dtype=np.uint16): + scale = log_scale_with_zero(range, n=n, dtype=x.dtype) + y = np.empty_like(x, dtype=dtype) + log_quant_with_zeros(x, y, np.log10(scale[1:])) + return (y, scale) + + +# optimized helper function for the above +@njit(parallel=True) +def log_quant_with_zeros(x, y, scale): + x = x.ravel() + y = y.ravel() + min_val = 10**scale[0] + + for i in prange(x.shape[0]): + # map small values to 0 + if x[i] < min_val: + y[i] = 0 + continue + + lx = np.log10(x[i]) + if lx >= scale[-1]: + # map too big values to max of scale + y[i] = len(scale) + else: + # binary search for the rest + k0 = 0 + k1 = len(scale) + while k1-k0 > 1: + km = k0 + (k1-k0)//2 + if lx < scale[km]: + k1 = km + else: + k0 = km + + if k0 == len(scale)-1: + q = k0 + elif k0 == 0: + q = 0 + else: + d0 = abs(lx-scale[k0]) + d1 = abs(lx-scale[k1]) + if d0 < d1: + q = k0 + else: + q = k1 + + y[i] = q+1 # add 1 to leave space for zero + + +@njit(parallel=True) +def average_pool(x, factor=2, missing=65535): + y = np.empty((x.shape[0]//factor, x.shape[1]//factor), dtype=x.dtype) + N = factor**2 + N_thresh = N//2 + + for iy in prange(y.shape[0]): + ix0 = iy * factor + ix1 = ix0 + factor + for jy in range(y.shape[1]): + jx0 = jy * factor + jx1 = jx0 + factor + v = float(0.0) + num_valid = 0 + + for ix in range(ix0, ix1): + for jx in range(jx0, jx1): + if x[ix,jx] != missing: + v += x[ix,jx] + num_valid += 1 + + if num_valid >= N_thresh: + y[iy,jy] = v/num_valid + else: + y[iy,jy] = missing + + return y + + +@njit(parallel=True) +def mode_pool(x, num_values=256, factor=2): + y = np.empty((x.shape[0]//factor, x.shape[1]//factor), dtype=x.dtype) + + for iy in prange(y.shape[0]): + v = np.empty(num_values, dtype=np.int64) + ix0 = iy * factor + ix1 = ix0 + factor + for jy in range(y.shape[1]): + jx0 = jy * factor + jx1 = jx0 + factor + v[:] = 0 + + for ix in range(ix0, ix1): + for jx in range(jx0, jx1): + v[x[ix,jx]] += 1 + + y[iy,jy] = v.argmax() + + return y + + +def fill_holes(missing=65535, rad=1): + def fill(x): + # identify mask of points to fill + o = np.ones((2*rad+1,2*rad+1), dtype=np.uint16) + n = np.prod(o.shape) + valid = (x != missing) + num_valid_neighbors = convolve(valid, o, mode='same', method='direct') + mask = ~valid & (num_valid_neighbors > 0) + + # compute mean of valid points around each fillable point + fx = x.copy() + fx[~valid] = 0 + mx = convolve(fx, o.astype(np.float64), mode='same', method='direct') + mx = mx[mask] / num_valid_neighbors[mask] + if np.issubdtype(x.dtype, np.integer): + mx = mx.round().astype(x.dtype) + + # fill holes with mean + fx = x.copy() + fx[mask] = mx + return fx + + return fill + diff --git a/ldcast/forecast.py b/ldcast/forecast.py new file mode 100644 index 0000000000000000000000000000000000000000..333e06616a33233fec2606368e79b72667d837fa --- /dev/null +++ b/ldcast/forecast.py @@ -0,0 +1,264 @@ +import contextlib +import gc + +import numpy as np +import torch +import torch.multiprocessing as mp + +from .features.transform import Antialiasing +from .models.autoenc import autoenc, encoder +from .models.genforecast import analysis, unet +from .models.diffusion import diffusion, plms + + +class Forecast: + def __init__( + self, + *, + ldm_weights_fn, + autoenc_weights_fn, + gpu='auto', + past_timesteps=4, + future_timesteps=20, + autoenc_time_ratio=4, + autoenc_hidden_dim=32, + verbose=True, + R_min_value=0.1, + R_zero_value=0.02, + R_min_output=0.1, + R_max_output=118.428, + log_R_mean=-0.051, + log_R_std=0.528, + ): + self.ldm_weights_fn = ldm_weights_fn + self.autoenc_weights_fn = autoenc_weights_fn + self.verbose = verbose + self.R_min_value = R_min_value + self.R_zero_value = R_zero_value + self.R_min_output = R_min_output + self.R_max_output = R_max_output + self.log_R_mean = log_R_mean + self.log_R_std = log_R_std + self.past_timesteps = past_timesteps + self.future_timesteps = future_timesteps + self.autoenc_time_ratio = autoenc_time_ratio + self.autoenc_hidden_dim = autoenc_hidden_dim + self.antialiasing = Antialiasing() + + # setup LDM + self.ldm = self._init_model() + if gpu is not None: + if gpu == 'auto': + if torch.cuda.device_count() > 0: + self.ldm.to(device="cuda") + else: + self.ldm.to(device=f"cuda:{gpu}") + # setup sampler + self.sampler = plms.PLMSSampler(self.ldm) + print(self.ldm.device) + gc.collect() + + def _init_model(self): + # setup autoencoder + enc = encoder.SimpleConvEncoder() + dec = encoder.SimpleConvDecoder() + autoencoder_obs = autoenc.AutoencoderKL(enc, dec) + #print(torch.load(self.autoenc_weights_fn)['state_dict'].keys()) + # autoencoder_obs.load_state_dict(torch.load(self.autoenc_weights_fn)['state_dict']) + autoencoder_obs.load_state_dict(torch.load(self.autoenc_weights_fn)) + autoencoders = [autoencoder_obs] + input_patches = [self.past_timesteps//self.autoenc_time_ratio] + input_size_ratios = [1] + embed_dim = [128] + analysis_depth = [4] + + # setup forecaster + analysis_net = analysis.AFNONowcastNetCascade( + autoencoders, + input_patches=input_patches, + input_size_ratios=input_size_ratios, + train_autoenc=False, + output_patches=self.future_timesteps//self.autoenc_time_ratio, + cascade_depth=3, + embed_dim=embed_dim, + analysis_depth=analysis_depth + ) + + # setup denoiser + denoiser = unet.UNetModel(in_channels=autoencoder_obs.hidden_width, + model_channels=256, out_channels=autoencoder_obs.hidden_width, + num_res_blocks=2, attention_resolutions=(1,2), + dims=3, channel_mult=(1, 2, 4), num_heads=8, + num_timesteps=self.future_timesteps//self.autoenc_time_ratio, + context_ch=analysis_net.cascade_dims + ) + + # create LDM + ldm = diffusion.LatentDiffusion(denoiser, autoencoder_obs, + context_encoder=analysis_net) + # ldm.load_state_dict(torch.load(self.ldm_weights_fn)['state_dict']) + ldm.load_state_dict(torch.load(self.ldm_weights_fn)) + return ldm + + def __call__( + self, + R_past, + num_diffusion_iters=50 + ): + # preprocess inputs and setup correct input shape + x = self.transform_precip(R_past) + timesteps = self.input_timesteps(x) + future_patches = self.future_timesteps // self.autoenc_time_ratio + gen_shape = (self.autoenc_hidden_dim, future_patches) + \ + (x.shape[-2]//4, x.shape[-1]//4) + x = [[x, timesteps]] + + # run LDM sampler + with contextlib.redirect_stdout(None): + (s, intermediates) = self.sampler.sample( + num_diffusion_iters, + x[0][0].shape[0], + gen_shape, + x, + progbar=self.verbose + ) + + # postprocess outputs + y_pred = self.ldm.autoencoder.decode(s) + R_pred = self.inv_transform_precip(y_pred) + + return R_pred[0,...] + + def transform_precip(self, R): + # x = R.copy() + x = R.clone().detach() + x[~(x >= self.R_min_value)] = self.R_zero_value + x = np.log10(x) + x -= self.log_R_mean + x /= self.log_R_std + x = x.reshape((1,) + x.shape) + x = self.antialiasing(x) + x = x.reshape((1,) + x.shape) + return torch.Tensor(x).to(device=self.ldm.device) + + def inv_transform_precip(self, x): + x *= self.log_R_std + x += self.log_R_mean + R = torch.pow(10, x) + if self.R_min_output: + R[R < self.R_min_output] = 0.0 + if self.R_max_output is not None: + R[R > self.R_max_output] = self.R_max_output + R = R[:,0,...] + return R.to(device='cpu').numpy() + + def input_timesteps(self, x): + batch_size = x.shape[0] + t0 = -x.shape[2]+1 + t1 = 1 + timesteps = torch.arange(t0, t1, + dtype=x.dtype, device=self.ldm.device) + return timesteps.unsqueeze(0).expand(batch_size,-1) + + +class ForecastDistributed: + def __init__( + self, + ldm_weights_fn, + autoenc_weights_fn, + past_timesteps=4, + future_timesteps=8, + autoenc_time_ratio=4, + autoenc_hidden_dim=32, + verbose=True, + R_min_value=0.1, + R_zero_value=0.02, + R_min_output=0.1, + R_max_output=118.428, + log_R_mean=-0.051, + log_R_std=0.528, + ): + self.verbose = verbose + self.R_min_value = R_min_value + self.R_zero_value = R_zero_value + self.R_min_output = R_min_output + self.R_max_output = R_max_output + self.log_R_mean = log_R_mean + self.log_R_std = log_R_std + self.past_timesteps = past_timesteps + self.future_timesteps = future_timesteps + self.autoenc_time_ratio = autoenc_time_ratio + self.autoenc_hidden_dim = autoenc_hidden_dim + + # start worker processes + context = mp.get_context('spawn') + self.input_queue = context.Queue() + self.output_queue = context.Queue() + process_kwargs = { + "past_timesteps": past_timesteps, + "future_timesteps": future_timesteps, + "ldm_weights_fn": ldm_weights_fn, + "autoenc_weights_fn": autoenc_weights_fn, + "autoenc_time_ratio": autoenc_time_ratio, + "autoenc_hidden_dim": autoenc_hidden_dim, + "R_min_value": R_min_value, + "R_zero_value": R_zero_value, + "R_min_output": R_min_output, + "R_max_output": R_max_output, + "log_R_mean": log_R_mean, + "log_R_std": log_R_std, + "verbose": True + } + self.num_procs = max(0, torch.cuda.device_count()) + self.compute_procs = mp.spawn( + _compute_process, + args=(self.input_queue, self.output_queue, process_kwargs), + nprocs=self.num_procs, + join=False + ) + + # wait for worker processes to be ready + for _ in range(self.num_procs): + self.output_queue.get() + + gc.collect() + + def __call__( + self, + R_past, + ensemble_members=1, + num_diffusion_iters=50 + ): + # send samples to compute processes + for (i, R_past_sample) in enumerate(R_past): + for j in range(ensemble_members): + self.input_queue.put((R_past_sample, num_diffusion_iters, i, j)) + + # build output array + pred_shape = (R_past.shape[0], self.future_timesteps) + \ + R_past.shape[2:] + (ensemble_members,) + R_pred = np.empty(pred_shape, R_past.dtype) + + # gather outputs from processes + predictions_needed = R_past.shape[0] * ensemble_members + for _ in range(predictions_needed): + (R_pred_sample, i, j) = self.output_queue.get() + R_pred[i,...,j] = R_pred_sample + + return R_pred + + def __del__(self): + for _ in range(self.num_procs): + self.input_queue.put(None) + self.compute_procs.join() + + +def _compute_process(process_index, input_queue, output_queue, kwargs): + gpu = process_index if (torch.cuda.device_count() > 0) else None + fc = Forecast(gpu=gpu, **kwargs) + output_queue.put("Ready") # signal process ready to accept inputs + + while (data := input_queue.get()) is not None: + (R_past, num_diffusion_iters, sample, member) = data + R_pred = fc(R_past, num_diffusion_iters=num_diffusion_iters) + output_queue.put((R_pred, sample, member)) diff --git a/ldcast/models/autoenc/autoenc.py b/ldcast/models/autoenc/autoenc.py new file mode 100644 index 0000000000000000000000000000000000000000..345e7e16e1b4dff5d2d1702ed2d449a0c5241842 --- /dev/null +++ b/ldcast/models/autoenc/autoenc.py @@ -0,0 +1,93 @@ +import torch +from torch import nn +import pytorch_lightning as pl + +from ..distributions import kl_from_standard_normal, ensemble_nll_normal +from ..distributions import sample_from_standard_normal + + +class AutoencoderKL(pl.LightningModule): + def __init__( + self, + encoder, decoder, + kl_weight=0.01, + encoded_channels=64, + hidden_width=32, + **kwargs + ): + super().__init__(**kwargs) + self.encoder = encoder + self.decoder = decoder + self.hidden_width = hidden_width + self.to_moments = nn.Conv3d(encoded_channels, 2*hidden_width, + kernel_size=1) + self.to_decoder = nn.Conv3d(hidden_width, encoded_channels, + kernel_size=1) + self.log_var = nn.Parameter(torch.zeros(size=())) + self.kl_weight = kl_weight + + def encode(self, x): + h = self.encoder(x) + (mean, log_var) = torch.chunk(self.to_moments(h), 2, dim=1) + return (mean, log_var) + + def decode(self, z): + z = self.to_decoder(z) + dec = self.decoder(z) + return dec + + def forward(self, input, sample_posterior=True): + (mean, log_var) = self.encode(input) + if sample_posterior: + z = sample_from_standard_normal(mean, log_var) + else: + z = mean + dec = self.decode(z) + return (dec, mean, log_var) + + def _loss(self, batch): + (x,y) = batch + while isinstance(x, list) or isinstance(x, tuple): + x = x[0][0] + (y_pred, mean, log_var) = self.forward(x) + + rec_loss = (y-y_pred).abs().mean() + kl_loss = kl_from_standard_normal(mean, log_var) + + total_loss = rec_loss + self.kl_weight * kl_loss + + return (total_loss, rec_loss, kl_loss) + + def training_step(self, batch, batch_idx): + loss = self._loss(batch)[0] + self.log("train_loss", loss) + return loss + + @torch.no_grad() + def val_test_step(self, batch, batch_idx, split="val"): + (total_loss, rec_loss, kl_loss) = self._loss(batch) + log_params = {"on_step": False, "on_epoch": True, "prog_bar": True} + self.log(f"{split}_loss", total_loss, **log_params) + self.log(f"{split}_rec_loss", rec_loss.mean(), **log_params) + self.log(f"{split}_kl_loss", kl_loss, **log_params) + + def validation_step(self, batch, batch_idx): + self.val_test_step(batch, batch_idx, split="val") + + def test_step(self, batch, batch_idx): + self.val_test_step(batch, batch_idx, split="test") + + def configure_optimizers(self): + optimizer = torch.optim.AdamW(self.parameters(), lr=1e-3, + betas=(0.5, 0.9), weight_decay=1e-3) + reduce_lr = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, patience=3, factor=0.25, verbose=True + ) + return { + "optimizer": optimizer, + "lr_scheduler": { + "scheduler": reduce_lr, + "monitor": "val_rec_loss", + "frequency": 1, + }, + } diff --git a/ldcast/models/autoenc/encoder.py b/ldcast/models/autoenc/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..eca28f1f6053d6f0a1269c44a7e2bd0b0f75e2d7 --- /dev/null +++ b/ldcast/models/autoenc/encoder.py @@ -0,0 +1,57 @@ +import numpy as np +import torch.nn as nn + +from ..blocks.resnet import ResBlock3D +from ..utils import activation, normalization + + +class SimpleConvEncoder(nn.Sequential): + def __init__(self, in_dim=1, levels=2, min_ch=64): + sequence = [] + channels = np.hstack([ + in_dim, + (8**np.arange(1,levels+1)).clip(min=min_ch) + ]) + + for i in range(levels): + in_channels = int(channels[i]) + out_channels = int(channels[i+1]) + res_kernel_size = (3,3,3) if i == 0 else (1,3,3) + res_block = ResBlock3D( + in_channels, out_channels, + kernel_size=res_kernel_size, + norm_kwargs={"num_groups": 1} + ) + sequence.append(res_block) + downsample = nn.Conv3d(out_channels, out_channels, + kernel_size=(2,2,2), stride=(2,2,2)) + sequence.append(downsample) + in_channels = out_channels + + super().__init__(*sequence) + + +class SimpleConvDecoder(nn.Sequential): + def __init__(self, in_dim=1, levels=2, min_ch=64): + sequence = [] + channels = np.hstack([ + in_dim, + (8**np.arange(1,levels+1)).clip(min=min_ch) + ]) + + for i in reversed(list(range(levels))): + in_channels = int(channels[i+1]) + out_channels = int(channels[i]) + upsample = nn.ConvTranspose3d(in_channels, in_channels, + kernel_size=(2,2,2), stride=(2,2,2)) + sequence.append(upsample) + res_kernel_size = (3,3,3) if (i == 0) else (1,3,3) + res_block = ResBlock3D( + in_channels, out_channels, + kernel_size=res_kernel_size, + norm_kwargs={"num_groups": 1} + ) + sequence.append(res_block) + in_channels = out_channels + + super().__init__(*sequence) diff --git a/ldcast/models/autoenc/training.py b/ldcast/models/autoenc/training.py new file mode 100644 index 0000000000000000000000000000000000000000..f4b53bbb2f4cafd257ed4d7d7cc412445cfc8e87 --- /dev/null +++ b/ldcast/models/autoenc/training.py @@ -0,0 +1,41 @@ +import pytorch_lightning as pl +import torch + +from . import autoenc + + +def setup_autoenc_training( + encoder, + decoder, + model_dir, +): + autoencoder = autoenc.AutoencoderKL(encoder, decoder) + + num_gpus = torch.cuda.device_count() + accelerator = "gpu" if (num_gpus > 0) else "cpu" + devices = torch.cuda.device_count() if (accelerator == "gpu") else 1 + + early_stopping = pl.callbacks.EarlyStopping( + "val_rec_loss", patience=6, verbose=True + ) + print(model_dir) + checkpoint = pl.callbacks.ModelCheckpoint( + dirpath=model_dir, + filename="{epoch}-{val_rec_loss:.4f}", + #filename=ckpt, + monitor="val_rec_loss", + every_n_epochs=1, + save_top_k=3, + save_weights_only=False, + ) + callbacks = [early_stopping, checkpoint] + + trainer = pl.Trainer( + accelerator=accelerator, + devices=devices, + max_epochs=1000, + #strategy='ddp' if (num_gpus > 1) else None, + callbacks=callbacks, + ) + + return (autoencoder, trainer) diff --git a/ldcast/models/benchmarks/dgmr.py b/ldcast/models/benchmarks/dgmr.py new file mode 100644 index 0000000000000000000000000000000000000000..704b9e46be0cd213037af7aed7576b2bda0c222a --- /dev/null +++ b/ldcast/models/benchmarks/dgmr.py @@ -0,0 +1,82 @@ +import gc + +import numpy as np +import tensorflow as tf + + +class DGMRModel: + def __init__( + self, + model_handle, + multi_gpu=True, + transform_to_rainrate=None, + transform_from_rainrate=None, + data_format='channels_first', + calibrated=False, + ): + self.transform_to_rainrate = transform_to_rainrate + self.transform_from_rainrate = transform_from_rainrate + self.data_format = data_format + self.calibrated = calibrated + + if multi_gpu and len(tf.config.list_physical_devices('GPU')) > 1: + # initialize multi-GPU strategy + strategy = tf.distribute.MirroredStrategy() + else: # use default strategy + strategy = tf.distribute.get_strategy() + + with strategy.scope(): + module = tf.saved_model.load(model_handle) + + self.model = module.signatures['default'] + input_signature = self.model.structured_input_signature[1] + self.noise_dim = input_signature['z'].shape[1] + self.past_timesteps = input_signature['labels$cond_frames'].shape[1] + + def __call__(self, x): + while isinstance(x, list) or isinstance(x, tuple): + x = x[0] + x = np.array(x, copy=False) + if self.data_format == "channels_first": + x = x.transpose(0,2,3,4,1) + if self.transform_to_rainrate is not None: + x = self.transform_to_rainrate(x) + x = tf.convert_to_tensor(x) + + num_samples = x.shape[0] + z = tf.random.normal(shape=(num_samples, self.noise_dim)) + if self.calibrated: + z = z * 2.0 + + onehot = tf.ones(shape=(num_samples, 1)) + inputs = { + "z": z, + "labels$onehot" : onehot, + "labels$cond_frames" : x + } + y = self.model(**inputs)['default'] + y = y[:,self.past_timesteps:,...] + + y = np.array(y) + if self.transform_from_rainrate is not None: + y = self.transform_from_rainrate(y) + if self.data_format == "channels_first": + y = y.transpose(0,4,1,2,3) + + return y + + +def create_ensemble( + dgmr, x, + ensemble_size=32, + model_path="../models/dgmr/256x256", + +): + y_pred = [] + for member in range(ensemble_size): + print(f"Generating member {member+1}/{ensemble_size}") + y_pred.append(dgmr(x)) + gc.collect() + + y_pred = np.stack(y_pred, axis=-1) + return y_pred diff --git a/ldcast/models/benchmarks/pysteps.py b/ldcast/models/benchmarks/pysteps.py new file mode 100644 index 0000000000000000000000000000000000000000..065af15d43993de9bd69951d9bdae4f9937d91aa --- /dev/null +++ b/ldcast/models/benchmarks/pysteps.py @@ -0,0 +1,106 @@ +# following https://pysteps.readthedocs.io/en/stable/auto_examples/plot_steps_nowcast.html +from datetime import timedelta + +import dask +import numpy as np +from pysteps import nowcasts +from pysteps.motion.lucaskanade import dense_lucaskanade +from pysteps.utils import transformation + + +class PySTEPSModel: + def __init__( + self, + data_format='channels_first', + future_timesteps=20, + ensemble_size=32, + km_per_pixel=1.0, + interval=timedelta(minutes=5), + transform_to_rainrate=None, + transform_from_rainrate=None, + ): + self.transform_to_rainrate = transform_to_rainrate + self.transform_from_rainrate = transform_from_rainrate + self.data_format = data_format + self.nowcast_method = nowcasts.get_method("steps") + self.future_timesteps = future_timesteps + self.ensemble_size = ensemble_size + self.km_per_pixel = km_per_pixel + self.interval = interval + + def zero_prediction(self, R, zerovalue): + out_shape = (self.future_timesteps,) + R.shape[1:] + \ + (self.ensemble_size,) + return np.full(out_shape, zerovalue, dtype=R.dtype) + + def predict_sample(self, x, threshold=-10.0, zerovalue=-15.0): + R = self.transform_to_rainrate(x) + (R, _) = transformation.dB_transform( + R, threshold=0.1, zerovalue=zerovalue + ) + R[~np.isfinite(R)] = zerovalue + if (R == zerovalue).all(): + R_f = self.zero_prediction(R, zerovalue) + else: + V = dense_lucaskanade(R) + try: + R_f = self.nowcast_method( + R, + V, + self.future_timesteps, + n_ens_members=self.ensemble_size, + n_cascade_levels=6, + precip_thr=threshold, + kmperpixel=self.km_per_pixel, + timestep=self.interval.total_seconds()/60, + noise_method="nonparametric", + vel_pert_method="bps", + mask_method="incremental", + num_workers=2 + ) + R_f = R_f.transpose(1,2,3,0) + except (ValueError, RuntimeError) as e: + zero_error = str(e).endswith("contains non-finite values") or \ + str(e).startswith("zero-size array to reduction operation") or \ + str(e).endswith("nonstationary AR(p) process") + if zero_error: + # occasional PySTEPS errors that happen with little/no precip + # therefore returning all zeros makes sense + R_f = self.zero_prediction(R, zerovalue) + else: + raise + + # Back-transform to rain rates + R_f = transformation.dB_transform( + R_f, threshold=threshold, inverse=True + )[0] + + if self.transform_from_rainrate is not None: + R_f = self.transform_from_rainrate(R_f) + + return R_f + + def __call__(self, x, parallel=True): + while isinstance(x, list) or isinstance(x, tuple): + x = x[0] + x = np.array(x, copy=False) + if self.data_format == "channels_first": + x = x.transpose(0,2,3,4,1) + + pred = self.predict_sample + if parallel: + pred = dask.delayed(pred) + y = [ + pred(x[i,:,:,:,0]) + for i in range(x.shape[0]) + ] + if parallel: + y = dask.compute(y, scheduler="threads", num_workers=len(y))[0] + y = np.stack(y, axis=0) + + if self.data_format == "channels_first": + y = np.expand_dims(y, 1) + else: + y = np.expand_dims(y, -2) + + return y diff --git a/ldcast/models/benchmarks/transform.py b/ldcast/models/benchmarks/transform.py new file mode 100644 index 0000000000000000000000000000000000000000..783b991220a945e42b5fae171cf7c9f4e50d00fb --- /dev/null +++ b/ldcast/models/benchmarks/transform.py @@ -0,0 +1,17 @@ +import numpy as np + + +def transform_to_rainrate(x, mean=-0.051, std=0.528, threshold=0.1): + x = x*std + mean + R = 10**x + R[R < threshold] = 0 + return R + + +def transform_from_rainrate( + R, mean=-0.051, std=0.528, + threshold=0.1, fill_value=0.02 +): + R = R.copy() + R[R < threshold] = fill_value + return (np.log10(R)-mean) / std diff --git a/ldcast/models/blocks/afno.py b/ldcast/models/blocks/afno.py new file mode 100644 index 0000000000000000000000000000000000000000..8bea2880bf6acf1c84dfdefc533c16f0ae6e55d5 --- /dev/null +++ b/ldcast/models/blocks/afno.py @@ -0,0 +1,348 @@ +#reference: https://github.com/NVlabs/AFNO-transformer +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from .attention import TemporalAttention + +class Mlp(nn.Module): + def __init__( + self, + in_features, hidden_features=None, out_features=None, + act_layer=nn.GELU, drop=0.0 + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) if drop > 0 else nn.Identity() + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class AFNO2D(nn.Module): + def __init__(self, hidden_size, num_blocks=8, sparsity_threshold=0.01, hard_thresholding_fraction=1, hidden_size_factor=1): + super().__init__() + assert hidden_size % num_blocks == 0, f"hidden_size {hidden_size} should be divisble by num_blocks {num_blocks}" + + self.hidden_size = hidden_size + self.sparsity_threshold = sparsity_threshold + self.num_blocks = num_blocks + self.block_size = self.hidden_size // self.num_blocks + self.hard_thresholding_fraction = hard_thresholding_fraction + self.hidden_size_factor = hidden_size_factor + self.scale = 0.02 + + self.w1 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size, self.block_size * self.hidden_size_factor)) + self.b1 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size * self.hidden_size_factor)) + self.w2 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size * self.hidden_size_factor, self.block_size)) + self.b2 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size)) + + def forward(self, x): + bias = x + + dtype = x.dtype + x = x.float() + B, H, W, C = x.shape + + x = torch.fft.rfft2(x, dim=(1, 2), norm="ortho") + x = x.reshape(B, H, W // 2 + 1, self.num_blocks, self.block_size) + + o1_real = torch.zeros([B, H, W // 2 + 1, self.num_blocks, self.block_size * self.hidden_size_factor], device=x.device) + o1_imag = torch.zeros([B, H, W // 2 + 1, self.num_blocks, self.block_size * self.hidden_size_factor], device=x.device) + o2_real = torch.zeros(x.shape, device=x.device) + o2_imag = torch.zeros(x.shape, device=x.device) + + total_modes = H // 2 + 1 + kept_modes = int(total_modes * self.hard_thresholding_fraction) + + o1_real[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes] = F.relu( + torch.einsum('...bi,bio->...bo', x[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes].real, self.w1[0]) - \ + torch.einsum('...bi,bio->...bo', x[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes].imag, self.w1[1]) + \ + self.b1[0] + ) + + o1_imag[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes] = F.relu( + torch.einsum('...bi,bio->...bo', x[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes].imag, self.w1[0]) + \ + torch.einsum('...bi,bio->...bo', x[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes].real, self.w1[1]) + \ + self.b1[1] + ) + + o2_real[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes] = ( + torch.einsum('...bi,bio->...bo', o1_real[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes], self.w2[0]) - \ + torch.einsum('...bi,bio->...bo', o1_imag[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes], self.w2[1]) + \ + self.b2[0] + ) + + o2_imag[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes] = ( + torch.einsum('...bi,bio->...bo', o1_imag[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes], self.w2[0]) + \ + torch.einsum('...bi,bio->...bo', o1_real[:, total_modes-kept_modes:total_modes+kept_modes, :kept_modes], self.w2[1]) + \ + self.b2[1] + ) + + x = torch.stack([o2_real, o2_imag], dim=-1) + x = F.softshrink(x, lambd=self.sparsity_threshold) + x = torch.view_as_complex(x) + x = x.reshape(B, H, W // 2 + 1, C) + x = torch.fft.irfft2(x, s=(H, W), dim=(1,2), norm="ortho") + x = x.type(dtype) + + return x + bias + + +class Block(nn.Module): + def __init__( + self, + dim, + mlp_ratio=4., + drop=0., + drop_path=0., + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + double_skip=True, + num_blocks=8, + sparsity_threshold=0.01, + hard_thresholding_fraction=1.0 + ): + super().__init__() + self.norm1 = norm_layer(dim) + self.filter = AFNO2D(dim, num_blocks, sparsity_threshold, hard_thresholding_fraction) + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + self.double_skip = double_skip + + def forward(self, x): + residual = x + x = self.norm1(x) + x = self.filter(x) + + if self.double_skip: + x = x + residual + residual = x + + x = self.norm2(x) + x = self.mlp(x) + x = x + residual + return x + + + +class AFNO3D(nn.Module): + def __init__( + self, hidden_size, num_blocks=8, sparsity_threshold=0.01, + hard_thresholding_fraction=1, hidden_size_factor=1 + ): + super().__init__() + assert hidden_size % num_blocks == 0, f"hidden_size {hidden_size} should be divisble by num_blocks {num_blocks}" + + self.hidden_size = hidden_size + self.sparsity_threshold = sparsity_threshold + self.num_blocks = num_blocks + self.block_size = self.hidden_size // self.num_blocks + self.hard_thresholding_fraction = hard_thresholding_fraction + self.hidden_size_factor = hidden_size_factor + self.scale = 0.02 + + self.w1 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size, self.block_size * self.hidden_size_factor)) + self.b1 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size * self.hidden_size_factor)) + self.w2 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size * self.hidden_size_factor, self.block_size)) + self.b2 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size)) + + def forward(self, x): + bias = x + + dtype = x.dtype + x = x.float() + B, D, H, W, C = x.shape + + x = torch.fft.rfftn(x, dim=(1, 2, 3), norm="ortho") + x = x.reshape(B, D, H, W // 2 + 1, self.num_blocks, self.block_size) + + o1_real = torch.zeros([B, D, H, W // 2 + 1, self.num_blocks, self.block_size * self.hidden_size_factor], device=x.device) + o1_imag = torch.zeros([B, D, H, W // 2 + 1, self.num_blocks, self.block_size * self.hidden_size_factor], device=x.device) + o2_real = torch.zeros(x.shape, device=x.device) + o2_imag = torch.zeros(x.shape, device=x.device) + + total_modes = H // 2 + 1 + kept_modes = int(total_modes * self.hard_thresholding_fraction) + + o1_real[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes] = F.relu( + torch.einsum('...bi,bio->...bo', x[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes].real, self.w1[0]) - \ + torch.einsum('...bi,bio->...bo', x[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes].imag, self.w1[1]) + \ + self.b1[0] + ) + + o1_imag[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes] = F.relu( + torch.einsum('...bi,bio->...bo', x[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes].imag, self.w1[0]) + \ + torch.einsum('...bi,bio->...bo', x[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes].real, self.w1[1]) + \ + self.b1[1] + ) + + o2_real[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes] = ( + torch.einsum('...bi,bio->...bo', o1_real[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes], self.w2[0]) - \ + torch.einsum('...bi,bio->...bo', o1_imag[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes], self.w2[1]) + \ + self.b2[0] + ) + + o2_imag[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes] = ( + torch.einsum('...bi,bio->...bo', o1_imag[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes], self.w2[0]) + \ + torch.einsum('...bi,bio->...bo', o1_real[:, :, total_modes-kept_modes:total_modes+kept_modes, :kept_modes], self.w2[1]) + \ + self.b2[1] + ) + + x = torch.stack([o2_real, o2_imag], dim=-1) + x = F.softshrink(x, lambd=self.sparsity_threshold) + x = torch.view_as_complex(x) + x = x.reshape(B, D, H, W // 2 + 1, C) + x = torch.fft.irfftn(x, s=(D, H, W), dim=(1,2,3), norm="ortho") + x = x.type(dtype) + + return x + bias + + +class AFNOBlock3d(nn.Module): + def __init__( + self, + dim, + mlp_ratio=4., + drop=0., + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + double_skip=True, + num_blocks=8, + sparsity_threshold=0.01, + hard_thresholding_fraction=1.0, + data_format="channels_last", + mlp_out_features=None, + ): + super().__init__() + self.norm_layer = norm_layer + self.norm1 = norm_layer(dim) + self.filter = AFNO3D(dim, num_blocks, sparsity_threshold, + hard_thresholding_fraction) + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp( + in_features=dim, out_features=mlp_out_features, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, drop=drop + ) + self.double_skip = double_skip + self.channels_first = (data_format == "channels_first") + + def forward(self, x): + if self.channels_first: + # AFNO natively uses a channels-last data format + x = x.permute(0,2,3,4,1) + + residual = x + x = self.norm1(x) + x = self.filter(x) + + if self.double_skip: + x = x + residual + residual = x + + x = self.norm2(x) + x = self.mlp(x) + x = x + residual + + if self.channels_first: + x = x.permute(0,4,1,2,3) + + return x + + +class PatchEmbed3d(nn.Module): + def __init__(self, patch_size=(4,4,4), in_chans=1, embed_dim=256): + super().__init__() + self.patch_size = patch_size + self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + + def forward(self, x): + x = self.proj(x) + x = x.permute(0,2,3,4,1) # convert to BHWC + return x + + +class PatchExpand3d(nn.Module): + def __init__(self, patch_size=(4,4,4), out_chans=1, embed_dim=256): + super().__init__() + self.patch_size = patch_size + self.proj = nn.Linear(embed_dim, out_chans*np.prod(patch_size)) + + def forward(self, x): + x = self.proj(x) + x = rearrange( + x, + "b d h w (p0 p1 p2 c_out) -> b c_out (d p0) (h p1) (w p2)", + p0=self.patch_size[0], + p1=self.patch_size[1], + p2=self.patch_size[2], + d=x.shape[1], + h=x.shape[2], + w=x.shape[3], + ) + return x + + +class AFNOCrossAttentionBlock3d(nn.Module): + """ AFNO 3D Block with channel mixing from two sources. + """ + def __init__( + self, + dim, + context_dim, + mlp_ratio=2., + drop=0., + act_layer=nn.GELU, + norm_layer=nn.Identity, + double_skip=True, + num_blocks=8, + sparsity_threshold=0.01, + hard_thresholding_fraction=1.0, + data_format="channels_last", + timesteps=None + ): + super().__init__() + + self.norm1 = norm_layer(dim) + self.norm2 = norm_layer(dim+context_dim) + mlp_hidden_dim = int((dim+context_dim) * mlp_ratio) + self.pre_proj = nn.Linear(dim+context_dim, dim+context_dim) + self.filter = AFNO3D(dim+context_dim, num_blocks, sparsity_threshold, + hard_thresholding_fraction) + self.mlp = Mlp( + in_features=dim+context_dim, + out_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, drop=drop + ) + self.channels_first = (data_format == "channels_first") + + def forward(self, x, y): + if self.channels_first: + # AFNO natively uses a channels-last order + x = x.permute(0,2,3,4,1) + y = y.permute(0,2,3,4,1) + + xy = torch.concat((self.norm1(x),y), axis=-1) + xy = self.pre_proj(xy) + xy + xy = self.filter(self.norm2(xy)) + xy # AFNO filter + x = self.mlp(xy) + x # feed-forward + + if self.channels_first: + x = x.permute(0,4,1,2,3) + + return x diff --git a/ldcast/models/blocks/attention.py b/ldcast/models/blocks/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..edcf9d716ace871d07b60ce375dc5e7ace4120cb --- /dev/null +++ b/ldcast/models/blocks/attention.py @@ -0,0 +1,104 @@ +import math + +import torch +from torch import nn +import torch.nn.functional as F + + +class TemporalAttention(nn.Module): + def __init__( + self, channels, context_channels=None, + head_dim=32, num_heads=8 + ): + super().__init__() + self.channels = channels + if context_channels is None: + context_channels = channels + self.context_channels = context_channels + self.head_dim = head_dim + self.num_heads = num_heads + self.inner_dim = head_dim * num_heads + self.attn_scale = self.head_dim ** -0.5 + if channels % num_heads: + raise ValueError("channels must be divisible by num_heads") + self.KV = nn.Linear(context_channels, self.inner_dim*2) + self.Q = nn.Linear(channels, self.inner_dim) + self.proj = nn.Linear(self.inner_dim, channels) + + def forward(self, x, y=None): + if y is None: + y = x + + (K,V) = self.KV(y).chunk(2, dim=-1) + (B, Dk, H, W, C) = K.shape + shape = (B, Dk, H, W, self.num_heads, self.head_dim) + K = K.reshape(shape) + V = V.reshape(shape) + + Q = self.Q(x) + (B, Dq, H, W, C) = Q.shape + shape = (B, Dq, H, W, self.num_heads, self.head_dim) + Q = Q.reshape(shape) + + K = K.permute((0,2,3,4,5,1)) # K^T + V = V.permute((0,2,3,4,1,5)) + Q = Q.permute((0,2,3,4,1,5)) + + attn = torch.matmul(Q, K) * self.attn_scale + attn = F.softmax(attn, dim=-1) + y = torch.matmul(attn, V) + y = y.permute((0,4,1,2,3,5)) + y = y.reshape((B,Dq,H,W,C)) + y = self.proj(y) + return y + + +class TemporalTransformer(nn.Module): + def __init__(self, + channels, + mlp_dim_mul=1, + **kwargs + ): + super().__init__() + self.attn1 = TemporalAttention(channels, **kwargs) + self.attn2 = TemporalAttention(channels, **kwargs) + self.norm1 = nn.LayerNorm(channels) + self.norm2 = nn.LayerNorm(channels) + self.norm3 = nn.LayerNorm(channels) + self.mlp = MLP(channels, dim_mul=mlp_dim_mul) + + def forward(self, x, y): + x = self.attn1(self.norm1(x)) + x # self attention + x = self.attn2(self.norm2(x), y) + x # cross attention + return self.mlp(self.norm3(x)) + x # feed-forward + + +class MLP(nn.Sequential): + def __init__(self, dim, dim_mul=4): + inner_dim = dim * dim_mul + sequence = [ + nn.Linear(dim, inner_dim), + nn.SiLU(), + nn.Linear(inner_dim, dim) + ] + super().__init__(*sequence) + + +def positional_encoding(position, dims, add_dims=()): + div_term = torch.exp( + torch.arange(0, dims, 2, device=position.device) * + (-math.log(10000.0) / dims) + ) + if position.ndim == 1: + arg = position[:,None] * div_term[None,:] + else: + arg = position[:,:,None] * div_term[None,None,:] + + pos_enc = torch.concat( + [torch.sin(arg), torch.cos(arg)], + dim=-1 + ) + if add_dims: + for dim in add_dims: + pos_enc = pos_enc.unsqueeze(dim) + return pos_enc diff --git a/ldcast/models/blocks/resnet.py b/ldcast/models/blocks/resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..b29bcb84815bb6ff7de48d2408e61ecfa0ca59b7 --- /dev/null +++ b/ldcast/models/blocks/resnet.py @@ -0,0 +1,70 @@ +from torch import nn +from torch.nn.utils.parametrizations import spectral_norm as sn + +from ..utils import activation, normalization + + +class ResBlock3D(nn.Module): + def __init__( + self, in_channels, out_channels, resample=None, + resample_factor=(1,1,1), kernel_size=(3,3,3), + act='swish', norm='group', norm_kwargs=None, + spectral_norm=False, + **kwargs + ): + super().__init__(**kwargs) + if in_channels != out_channels: + self.proj = nn.Conv3d(in_channels, out_channels, kernel_size=1) + else: + self.proj = nn.Identity() + + padding = tuple(k//2 for k in kernel_size) + if resample == "down": + self.resample = nn.AvgPool3d(resample_factor, ceil_mode=True) + self.conv1 = nn.Conv3d(in_channels, out_channels, + kernel_size=kernel_size, stride=resample_factor, padding=padding) + self.conv2 = nn.Conv3d(out_channels, out_channels, + kernel_size=kernel_size, padding=padding) + elif resample == "up": + self.resample = nn.Upsample( + scale_factor=resample_factor, mode='trilinear') + self.conv1 = nn.ConvTranspose3d(in_channels, out_channels, + kernel_size=kernel_size, padding=padding) + output_padding = tuple( + 2*p+s-k for (p,s,k) in zip(padding,resample_factor,kernel_size) + ) + self.conv2 = nn.ConvTranspose3d(out_channels, out_channels, + kernel_size=kernel_size, stride=resample_factor, + padding=padding, output_padding=output_padding) + else: + self.resample = nn.Identity() + self.conv1 = nn.Conv3d(in_channels, out_channels, + kernel_size=kernel_size, padding=padding) + self.conv2 = nn.Conv3d(out_channels, out_channels, + kernel_size=kernel_size, padding=padding) + + if isinstance(act, str): + act = (act, act) + self.act1 = activation(act_type=act[0]) + self.act2 = activation(act_type=act[1]) + + if norm_kwargs is None: + norm_kwargs = {} + self.norm1 = normalization(in_channels, norm_type=norm, **norm_kwargs) + self.norm2 = normalization(out_channels, norm_type=norm, **norm_kwargs) + if spectral_norm: + self.conv1 = sn(self.conv1) + self.conv2 = sn(self.conv2) + if not isinstance(self.proj, nn.Identity): + self.proj = sn(self.proj) + + + def forward(self, x): + x_in = self.resample(self.proj(x)) + x = self.norm1(x) + x = self.act1(x) + x = self.conv1(x) + x = self.norm2(x) + x = self.act2(x) + x = self.conv2(x) + return x + x_in diff --git a/ldcast/models/diffusion/diffusion.py b/ldcast/models/diffusion/diffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..9917fcfede4fc44aa76f1c9d4b16fcb508e6beb1 --- /dev/null +++ b/ldcast/models/diffusion/diffusion.py @@ -0,0 +1,222 @@ +""" +From https://github.com/CompVis/latent-diffusion/main/ldm/models/diffusion/ddpm.py +Pared down to simplify code. + +The original file acknowledges: +https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +https://github.com/openai/improved-diffusion/blob/e94489283bb876ac1477d5dd7709bbbd2d9902ce/improved_diffusion/gaussian_diffusion.py +https://github.com/CompVis/taming-transformers +""" + +import torch +import torch.nn as nn +import numpy as np +import pytorch_lightning as pl +from contextlib import contextmanager +from functools import partial +from torchmetrics import MeanSquaredError + +from .utils import make_beta_schedule, extract_into_tensor, noise_like, timestep_embedding +from .ema import LitEma +from ..blocks.afno import PatchEmbed3d, PatchExpand3d, AFNOBlock3d + + +class LatentDiffusion(pl.LightningModule): + def __init__(self, + model, + autoencoder, + context_encoder=None, + timesteps=1000, + beta_schedule="linear", + loss_type="l2", + use_ema=True, + lr=1e-4, + lr_warmup=0, + linear_start=1e-4, + linear_end=2e-2, + cosine_s=8e-3, + parameterization="eps", # all assuming fixed variance schedules + ): + super().__init__() + self.model = model + self.autoencoder = autoencoder.requires_grad_(False) + self.conditional = (context_encoder is not None) + self.context_encoder = context_encoder + self.lr = lr + self.lr_warmup = lr_warmup + + self.val_loss = MeanSquaredError() + + assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"' + self.parameterization = parameterization + + self.use_ema = use_ema + if self.use_ema: + self.model_ema = LitEma(self.model) + + self.register_schedule( + beta_schedule=beta_schedule, timesteps=timesteps, + linear_start=linear_start, linear_end=linear_end, + cosine_s=cosine_s + ) + + self.loss_type = loss_type + + def register_schedule(self, beta_schedule="linear", timesteps=1000, + linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): + + betas = make_beta_schedule( + beta_schedule, timesteps, + linear_start=linear_start, linear_end=linear_end, + cosine_s=cosine_s + ) + alphas = 1. - betas + alphas_cumprod = np.cumprod(alphas, axis=0) + alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) + + timesteps, = betas.shape + self.num_timesteps = int(timesteps) + self.linear_start = linear_start + self.linear_end = linear_end + assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep' + + to_torch = partial(torch.tensor, dtype=torch.float32) + + self.register_buffer('betas', to_torch(betas)) + self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) + self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) + + # calculations for diffusion q(x_t | x_{t-1}) and others + self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) + self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) + + @contextmanager + def ema_scope(self, context=None): + if self.use_ema: + self.model_ema.store(self.model.parameters()) + self.model_ema.copy_to(self.model) + if context is not None: + print(f"{context}: Switched to EMA weights") + try: + yield None + finally: + if self.use_ema: + self.model_ema.restore(self.model.parameters()) + if context is not None: + print(f"{context}: Restored training weights") + + def apply_model(self, x_noisy, t, cond=None, return_ids=False): + if self.conditional: + cond = self.context_encoder(cond) + with self.ema_scope(): + return self.model(x_noisy, t, context=cond) + + def q_sample(self, x_start, t, noise=None): + if noise is None: + noise = torch.randn_like(x_start) + return ( + extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise + ) + + def get_loss(self, pred, target, mean=True): + if self.loss_type == 'l1': + loss = (target - pred).abs() + if mean: + loss = loss.mean() + elif self.loss_type == 'l2': + if mean: + loss = torch.nn.functional.mse_loss(target, pred) + else: + loss = torch.nn.functional.mse_loss(target, pred, reduction='none') + else: + raise NotImplementedError("unknown loss type '{loss_type}'") + + return loss + + def p_losses(self, x_start, t, noise=None, context=None): + if noise is None: + noise = torch.randn_like(x_start) + x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) + model_out = self.model(x_noisy, t, context=context) + + if self.parameterization == "eps": + target = noise + elif self.parameterization == "x0": + target = x_start + else: + raise NotImplementedError(f"Parameterization {self.parameterization} not yet supported") + + return self.get_loss(model_out, target, mean=False).mean() + + def forward(self, x, *args, **kwargs): + t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long() + return self.p_losses(x, t, *args, **kwargs) + + def shared_step(self, batch): + (x,y) = batch + y = self.autoencoder.encode(y)[0] + context = self.context_encoder(x) if self.conditional else None + return self(y, context=context) + + def training_step(self, batch, batch_idx): + loss = self.shared_step(batch) + self.log("train_loss", loss) + return loss + + @torch.no_grad() + def validation_step(self, batch, batch_idx): + #x, y = batch + #y_pred = self(x) + #loss2 = torch.nn.functional.mse_loss(y_pred, y) + loss = self.shared_step(batch) + with self.ema_scope(): + loss_ema = self.shared_step(batch) + log_params = {"on_step": False, "on_epoch": True, "prog_bar": True} + self.log("val_loss", loss, **log_params) + self.log("val_loss_ema", loss, **log_params) + #self.log("mean_square_error", loss2, **log_params) + + def test_step(self, batch, batch_idx): + return self.validation_step(batch, batch_idx) + + def on_train_batch_end(self, *args, **kwargs): + if self.use_ema: + self.model_ema(self.model) + + def configure_optimizers(self): + optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr, + betas=(0.5, 0.9), weight_decay=1e-3) + reduce_lr = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, patience=3, factor=0.25, verbose=True + ) + return { + "optimizer": optimizer, + "lr_scheduler": { + "scheduler": reduce_lr, + "monitor": "val_loss_ema", + "frequency": 1, + }, + } + + def optimizer_step( + self, + epoch, + batch_idx, + optimizer, + optimizer_idx, + #optimizer_closure, + **kwargs + ): + if self.trainer.global_step < self.lr_warmup: + lr_scale = (self.trainer.global_step+1) / self.lr_warmup + for pg in optimizer.param_groups: + pg['lr'] = lr_scale * self.lr + + super().optimizer_step( + epoch, batch_idx, optimizer, + optimizer_idx, + #optimizer_closure, + **kwargs + ) + diff --git a/ldcast/models/diffusion/ema.py b/ldcast/models/diffusion/ema.py new file mode 100644 index 0000000000000000000000000000000000000000..102a4ba7908a3dda4897e7273ef977c1ba5629ac --- /dev/null +++ b/ldcast/models/diffusion/ema.py @@ -0,0 +1,76 @@ +import torch +from torch import nn + + +class LitEma(nn.Module): + def __init__(self, model, decay=0.9999, use_num_upates=True): + super().__init__() + if decay < 0.0 or decay > 1.0: + raise ValueError('Decay must be between 0 and 1') + + self.m_name2s_name = {} + self.register_buffer('decay', torch.tensor(decay, dtype=torch.float32)) + self.register_buffer('num_updates', torch.tensor(0,dtype=torch.int) if use_num_upates + else torch.tensor(-1,dtype=torch.int)) + + for name, p in model.named_parameters(): + if p.requires_grad: + #remove as '.'-character is not allowed in buffers + s_name = name.replace('.','') + self.m_name2s_name.update({name:s_name}) + self.register_buffer(s_name,p.clone().detach().data) + + self.collected_params = [] + + def forward(self,model): + decay = self.decay + + if self.num_updates >= 0: + self.num_updates += 1 + decay = min(self.decay,(1 + self.num_updates) / (10 + self.num_updates)) + + one_minus_decay = 1.0 - decay + + with torch.no_grad(): + m_param = dict(model.named_parameters()) + shadow_params = dict(self.named_buffers()) + + for key in m_param: + if m_param[key].requires_grad: + sname = self.m_name2s_name[key] + shadow_params[sname] = shadow_params[sname].type_as(m_param[key]) + shadow_params[sname].sub_(one_minus_decay * (shadow_params[sname] - m_param[key])) + else: + assert not key in self.m_name2s_name + + def copy_to(self, model): + m_param = dict(model.named_parameters()) + shadow_params = dict(self.named_buffers()) + for key in m_param: + if m_param[key].requires_grad: + m_param[key].data.copy_(shadow_params[self.m_name2s_name[key]].data) + else: + assert not key in self.m_name2s_name + + def store(self, parameters): + """ + Save the current parameters for restoring later. + Args: + parameters: Iterable of `torch.nn.Parameter`; the parameters to be + temporarily stored. + """ + self.collected_params = [param.clone() for param in parameters] + + def restore(self, parameters): + """ + Restore the parameters stored with the `store` method. + Useful to validate the model with EMA parameters without affecting the + original optimization process. Store the parameters before the + `copy_to` method. After validation (or model saving), use this to + restore the former parameters. + Args: + parameters: Iterable of `torch.nn.Parameter`; the parameters to be + updated with the stored parameters. + """ + for c_param, param in zip(self.collected_params, parameters): + param.data.copy_(c_param.data) diff --git a/ldcast/models/diffusion/plms.py b/ldcast/models/diffusion/plms.py new file mode 100644 index 0000000000000000000000000000000000000000..beda5c80527b508f972dc1f898ab3fd7ce2eca49 --- /dev/null +++ b/ldcast/models/diffusion/plms.py @@ -0,0 +1,245 @@ +""" +From: https://github.com/CompVis/latent-diffusion/blob/main/ldm/models/diffusion/plms.py +""" + + +"""SAMPLING ONLY.""" + +import torch +import numpy as np +from tqdm import tqdm + +from .utils import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like + + +class PLMSSampler: + def __init__(self, model, schedule="linear", **kwargs): + self.model = model + self.ddpm_num_timesteps = model.num_timesteps + self.schedule = schedule + + def register_buffer(self, name, attr): + #if type(attr) == torch.Tensor: + # if attr.device != torch.device("cuda"): + # attr = attr.to(torch.device("cuda")) + setattr(self, name, attr) + + def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True): + if ddim_eta != 0: + raise ValueError('ddim_eta must be 0 for PLMS') + self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps, + num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose) + alphas_cumprod = self.model.alphas_cumprod + assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep' + to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device) + + self.register_buffer('betas', to_torch(self.model.betas)) + self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) + self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev)) + + # calculations for diffusion q(x_t | x_{t-1}) and others + self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu()))) + self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu()))) + self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu()))) + self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu()))) + self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1))) + + # ddim sampling parameters + ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(), + ddim_timesteps=self.ddim_timesteps, + eta=ddim_eta,verbose=verbose) + self.register_buffer('ddim_sigmas', ddim_sigmas) + self.register_buffer('ddim_alphas', ddim_alphas) + self.register_buffer('ddim_alphas_prev', ddim_alphas_prev) + self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas)) + sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt( + (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * ( + 1 - self.alphas_cumprod / self.alphas_cumprod_prev)) + self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps) + + @torch.no_grad() + def sample(self, + S, + batch_size, + shape, + conditioning=None, + callback=None, + normals_sequence=None, + img_callback=None, + quantize_x0=False, + eta=0., + mask=None, + x0=None, + temperature=1., + noise_dropout=0., + score_corrector=None, + corrector_kwargs=None, + verbose=True, + x_T=None, + log_every_t=100, + unconditional_guidance_scale=1., + unconditional_conditioning=None, + progbar=True, + # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ... + **kwargs + ): + """ + if conditioning is not None: + if isinstance(conditioning, dict): + cbs = conditioning[list(conditioning.keys())[0]].shape[0] + if cbs != batch_size: + print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}") + else: + if conditioning.shape[0] != batch_size: + print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}") + """ + + self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose) + # sampling + size = (batch_size,) + shape + print(f'Data shape for PLMS sampling is {size}') + + samples, intermediates = self.plms_sampling(conditioning, size, + callback=callback, + img_callback=img_callback, + quantize_denoised=quantize_x0, + mask=mask, x0=x0, + ddim_use_original_steps=False, + noise_dropout=noise_dropout, + temperature=temperature, + score_corrector=score_corrector, + corrector_kwargs=corrector_kwargs, + x_T=x_T, + log_every_t=log_every_t, + unconditional_guidance_scale=unconditional_guidance_scale, + unconditional_conditioning=unconditional_conditioning, + progbar=progbar + ) + return samples, intermediates + + @torch.no_grad() + def plms_sampling(self, cond, shape, + x_T=None, ddim_use_original_steps=False, + callback=None, timesteps=None, quantize_denoised=False, + mask=None, x0=None, img_callback=None, log_every_t=100, + temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None, + unconditional_guidance_scale=1., unconditional_conditioning=None, progbar=True): + device = self.model.betas.device + b = shape[0] + if x_T is None: + img = torch.randn(shape, device=device) + else: + img = x_T + + if timesteps is None: + timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps + elif timesteps is not None and not ddim_use_original_steps: + subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1 + timesteps = self.ddim_timesteps[:subset_end] + + intermediates = {'x_inter': [img], 'pred_x0': [img]} + time_range = list(reversed(range(0,timesteps))) if ddim_use_original_steps else np.flip(timesteps) + total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0] + print(f"Running PLMS Sampling with {total_steps} timesteps") + + iterator = time_range + if progbar: + iterator = tqdm(iterator, desc='PLMS Sampler', total=total_steps) + old_eps = [] + + for i, step in enumerate(iterator): + index = total_steps - i - 1 + ts = torch.full((b,), step, device=device, dtype=torch.long) + ts_next = torch.full((b,), time_range[min(i + 1, len(time_range) - 1)], device=device, dtype=torch.long) + + if mask is not None: + assert x0 is not None + img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass? + img = img_orig * mask + (1. - mask) * img + + outs = self.p_sample_plms(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps, + quantize_denoised=quantize_denoised, temperature=temperature, + noise_dropout=noise_dropout, score_corrector=score_corrector, + corrector_kwargs=corrector_kwargs, + unconditional_guidance_scale=unconditional_guidance_scale, + unconditional_conditioning=unconditional_conditioning, + old_eps=old_eps, t_next=ts_next) + img, pred_x0, e_t = outs + old_eps.append(e_t) + if len(old_eps) >= 4: + old_eps.pop(0) + if callback: callback(i) + if img_callback: img_callback(pred_x0, i) + + if index % log_every_t == 0 or index == total_steps - 1: + intermediates['x_inter'].append(img) + intermediates['pred_x0'].append(pred_x0) + + return img, intermediates + + @torch.no_grad() + def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False, + temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None, + unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None): + b, *_, device = *x.shape, x.device + + def get_model_output(x, t): + if unconditional_conditioning is None or unconditional_guidance_scale == 1.: + e_t = self.model.apply_model(x, t, c) + else: + x_in = torch.cat([x] * 2) + t_in = torch.cat([t] * 2) + c_in = torch.cat([unconditional_conditioning, c]) + e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2) + e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond) + + if score_corrector is not None: + assert self.model.parameterization == "eps" + e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs) + + return e_t + + alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas + alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev + sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas + sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas + + def get_x_prev_and_pred_x0(e_t, index): + # select parameters corresponding to the currently considered timestep + param_shape = (b,) + (1,)*(x.ndim-1) + a_t = torch.full(param_shape, alphas[index], device=device) + a_prev = torch.full(param_shape, alphas_prev[index], device=device) + sigma_t = torch.full(param_shape, sigmas[index], device=device) + sqrt_one_minus_at = torch.full(param_shape, sqrt_one_minus_alphas[index],device=device) + + # current prediction for x_0 + pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt() + if quantize_denoised: + pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0) + # direction pointing to x_t + dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t + noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature + if noise_dropout > 0.: + noise = torch.nn.functional.dropout(noise, p=noise_dropout) + x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise + return x_prev, pred_x0 + + e_t = get_model_output(x, t) + if len(old_eps) == 0: + # Pseudo Improved Euler (2nd order) + x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index) + e_t_next = get_model_output(x_prev, t_next) + e_t_prime = (e_t + e_t_next) / 2 + elif len(old_eps) == 1: + # 2nd order Pseudo Linear Multistep (Adams-Bashforth) + e_t_prime = (3 * e_t - old_eps[-1]) / 2 + elif len(old_eps) == 2: + # 3nd order Pseudo Linear Multistep (Adams-Bashforth) + e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12 + elif len(old_eps) >= 3: + # 4nd order Pseudo Linear Multistep (Adams-Bashforth) + e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24 + + x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index) + + return x_prev, pred_x0, e_t diff --git a/ldcast/models/diffusion/utils.py b/ldcast/models/diffusion/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bc21ead1dfccc512285ef0685e49c577e5250a51 --- /dev/null +++ b/ldcast/models/diffusion/utils.py @@ -0,0 +1,246 @@ +# adopted from +# https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py +# and +# https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +# and +# https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py +# +# thanks! + +import os +import math +import torch +import torch.nn as nn +import numpy as np +from einops import repeat + + +def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): + if schedule == "linear": + betas = ( + torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2 + ) + + elif schedule == "cosine": + timesteps = ( + torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s + ) + alphas = timesteps / (1 + cosine_s) * np.pi / 2 + alphas = torch.cos(alphas).pow(2) + alphas = alphas / alphas[0] + betas = 1 - alphas[1:] / alphas[:-1] + betas = np.clip(betas, a_min=0, a_max=0.999) + + elif schedule == "sqrt_linear": + betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) + elif schedule == "sqrt": + betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5 + else: + raise ValueError(f"schedule '{schedule}' unknown.") + return betas.numpy() + + +def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True): + if ddim_discr_method == 'uniform': + c = num_ddpm_timesteps // num_ddim_timesteps + ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c))) + elif ddim_discr_method == 'quad': + ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int) + else: + raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"') + + # assert ddim_timesteps.shape[0] == num_ddim_timesteps + # add one to get the final alpha values right (the ones from first scale to data during sampling) + steps_out = ddim_timesteps + 1 + if verbose: + print(f'Selected timesteps for ddim sampler: {steps_out}') + return steps_out + + +def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True): + # select alphas for computing the variance schedule + alphas = alphacums[ddim_timesteps] + alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist()) + + # according the the formula provided in https://arxiv.org/abs/2010.02502 + sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev)) + if verbose: + print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}') + print(f'For the chosen value of eta, which is {eta}, ' + f'this results in the following sigma_t schedule for ddim sampler {sigmas}') + return sigmas, alphas, alphas_prev + + +def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): + """ + Create a beta schedule that discretizes the given alpha_t_bar function, + which defines the cumulative product of (1-beta) over time from t = [0,1]. + :param num_diffusion_timesteps: the number of betas to produce. + :param alpha_bar: a lambda that takes an argument t from 0 to 1 and + produces the cumulative product of (1-beta) up to that + part of the diffusion process. + :param max_beta: the maximum beta to use; use values lower than 1 to + prevent singularities. + """ + betas = [] + for i in range(num_diffusion_timesteps): + t1 = i / num_diffusion_timesteps + t2 = (i + 1) / num_diffusion_timesteps + betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) + return np.array(betas) + + +def extract_into_tensor(a, t, x_shape): + b, *_ = t.shape + out = a.gather(-1, t) + return out.reshape(b, *((1,) * (len(x_shape) - 1))) + + +def checkpoint(func, inputs, params, flag): + """ + Evaluate a function without caching intermediate activations, allowing for + reduced memory at the expense of extra compute in the backward pass. + :param func: the function to evaluate. + :param inputs: the argument sequence to pass to `func`. + :param params: a sequence of parameters `func` depends on but does not + explicitly take as arguments. + :param flag: if False, disable gradient checkpointing. + """ + if flag: + args = tuple(inputs) + tuple(params) + return CheckpointFunction.apply(func, len(inputs), *args) + else: + return func(*inputs) + + +class CheckpointFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, run_function, length, *args): + ctx.run_function = run_function + ctx.input_tensors = list(args[:length]) + ctx.input_params = list(args[length:]) + + with torch.no_grad(): + output_tensors = ctx.run_function(*ctx.input_tensors) + return output_tensors + + @staticmethod + def backward(ctx, *output_grads): + ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] + with torch.enable_grad(): + # Fixes a bug where the first op in run_function modifies the + # Tensor storage in place, which is not allowed for detach()'d + # Tensors. + shallow_copies = [x.view_as(x) for x in ctx.input_tensors] + output_tensors = ctx.run_function(*shallow_copies) + input_grads = torch.autograd.grad( + output_tensors, + ctx.input_tensors + ctx.input_params, + output_grads, + allow_unused=True, + ) + del ctx.input_tensors + del ctx.input_params + del output_tensors + return (None, None) + input_grads + + +def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False): + """ + Create sinusoidal timestep embeddings. + :param timesteps: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an [N x dim] Tensor of positional embeddings. + """ + if not repeat_only: + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half + ).to(device=timesteps.device) + args = timesteps[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + else: + embedding = repeat(timesteps, 'b -> b d', d=dim) + return embedding + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module + + +def scale_module(module, scale): + """ + Scale the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().mul_(scale) + return module + + +def mean_flat(tensor): + """ + Take the mean over all non-batch dimensions. + """ + return tensor.mean(dim=list(range(1, len(tensor.shape)))) + + +class GroupNorm32(nn.GroupNorm): + def forward(self, x): + return super().forward(x.float()).type(x.dtype) + + +def normalization(channels): + """ + Make a standard normalization layer. + :param channels: number of input channels. + :return: an nn.Module for normalization. + """ + return nn.Identity() #GroupNorm32(32, channels) + + +def noise_like(shape, device, repeat=False): + repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1))) + noise = lambda: torch.randn(shape, device=device) + return repeat_noise() if repeat else noise() + + +def conv_nd(dims, *args, **kwargs): + """ + Create a 1D, 2D, or 3D convolution module. + """ + if dims == 1: + return nn.Conv1d(*args, **kwargs) + elif dims == 2: + return nn.Conv2d(*args, **kwargs) + elif dims == 3: + return nn.Conv3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +def linear(*args, **kwargs): + """ + Create a linear module. + """ + return nn.Linear(*args, **kwargs) + + +def avg_pool_nd(dims, *args, **kwargs): + """ + Create a 1D, 2D, or 3D average pooling module. + """ + if dims == 1: + return nn.AvgPool1d(*args, **kwargs) + elif dims == 2: + return nn.AvgPool2d(*args, **kwargs) + elif dims == 3: + return nn.AvgPool3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") \ No newline at end of file diff --git a/ldcast/models/distributions.py b/ldcast/models/distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..ade8711d8b4fabe3a2db75814fee5dcb8f894feb --- /dev/null +++ b/ldcast/models/distributions.py @@ -0,0 +1,29 @@ +import numpy as np +import torch + + +def kl_from_standard_normal(mean, log_var): + kl = 0.5 * (log_var.exp() + mean.square() - 1.0 - log_var) + return kl.mean() + + +def sample_from_standard_normal(mean, log_var, num=None): + std = (0.5 * log_var).exp() + shape = mean.shape + if num is not None: + # expand channel 1 to create several samples + shape = shape[:1] + (num,) + shape[1:] + mean = mean[:,None,...] + std = std[:,None,...] + return mean + std * torch.randn(shape, device=mean.device) + + +def ensemble_nll_normal(ensemble, sample, epsilon=1e-5): + mean = ensemble.mean(dim=1) + var = ensemble.var(dim=1, unbiased=True) + epsilon + logvar = var.log() + + diff = sample[:,None,...] - mean + logtwopi = np.log(2*np.pi) + nll = (logtwopi + logvar + diff.square() / var).mean() + return nll diff --git a/ldcast/models/genforecast/analysis.py b/ldcast/models/genforecast/analysis.py new file mode 100644 index 0000000000000000000000000000000000000000..7a423931b16290c49ce583bf13dc738b756e9ea8 --- /dev/null +++ b/ldcast/models/genforecast/analysis.py @@ -0,0 +1,33 @@ +import torch +from torch import nn +from torch.nn import functional as F + +from ..nowcast.nowcast import AFNONowcastNetBase +from ..blocks.resnet import ResBlock3D + + +class AFNONowcastNetCascade(AFNONowcastNetBase): + def __init__(self, *args, cascade_depth=4, **kwargs): + super().__init__(*args, **kwargs) + self.cascade_depth = cascade_depth + self.resnet = nn.ModuleList() + ch = self.embed_dim_out + self.cascade_dims = [ch] + for i in range(cascade_depth-1): + ch_out = 2*ch + self.cascade_dims.append(ch_out) + self.resnet.append( + ResBlock3D(ch, ch_out, kernel_size=(1,3,3), norm=None) + ) + ch = ch_out + + def forward(self, x): + x = super().forward(x) + img_shape = tuple(x.shape[-2:]) + cascade = {img_shape: x} + for i in range(self.cascade_depth-1): + x = F.avg_pool3d(x, (1,2,2)) + x = self.resnet[i](x) + img_shape = tuple(x.shape[-2:]) + cascade[img_shape] = x + return cascade diff --git a/ldcast/models/genforecast/training.py b/ldcast/models/genforecast/training.py new file mode 100644 index 0000000000000000000000000000000000000000..802fe5995844decf22ce506091464d5261a8c2b4 --- /dev/null +++ b/ldcast/models/genforecast/training.py @@ -0,0 +1,42 @@ +import pytorch_lightning as pl +import torch + +from ..diffusion import diffusion + + +def setup_genforecast_training( + model, + autoencoder, + context_encoder, + model_dir, + lr=1e-4 +): + ldm = diffusion.LatentDiffusion(model, autoencoder, + context_encoder=context_encoder, lr=lr) + + num_gpus = torch.cuda.device_count() + accelerator = "gpu" if (num_gpus > 0) else "cpu" + devices = torch.cuda.device_count() if (accelerator == "gpu") else 1 + + early_stopping = pl.callbacks.EarlyStopping( + "val_loss_ema", patience=6, verbose=True, check_finite=False + ) + checkpoint = pl.callbacks.ModelCheckpoint( + dirpath=model_dir, + filename="{epoch}-{val_loss_ema:.4f}", + monitor="val_loss_ema", + every_n_epochs=1, + save_top_k=3 + ) + callbacks = [early_stopping, checkpoint] + + trainer = pl.Trainer( + accelerator=accelerator, + devices=devices, + max_epochs=300, + #strategy='dp' if (num_gpus > 1) else None, + callbacks=callbacks, + #precision=16 + ) + + return (ldm, trainer) diff --git a/ldcast/models/genforecast/unet.py b/ldcast/models/genforecast/unet.py new file mode 100644 index 0000000000000000000000000000000000000000..4152b93864135f3cedefbee701a99fc51c318ac8 --- /dev/null +++ b/ldcast/models/genforecast/unet.py @@ -0,0 +1,489 @@ +from abc import abstractmethod +from functools import partial +import math +from typing import Iterable + +import numpy as np +import torch as th +import torch.nn as nn +import torch.nn.functional as F + +from ..diffusion.utils import ( + checkpoint, + conv_nd, + linear, + avg_pool_nd, + zero_module, + normalization, + timestep_embedding, +) +from ..blocks.afno import AFNOCrossAttentionBlock3d +SpatialTransformer = type(None) +#from ldm.modules.attention import SpatialTransformer + + +class TimestepBlock(nn.Module): + """ + Any module where forward() takes timestep embeddings as a second argument. + """ + + @abstractmethod + def forward(self, x, emb): + """ + Apply the module to `x` given `emb` timestep embeddings. + """ + + +class TimestepEmbedSequential(nn.Sequential, TimestepBlock): + """ + A sequential module that passes timestep embeddings to the children that + support it as an extra input. + """ + + def forward(self, x, emb, context=None): + for layer in self: + if isinstance(layer, TimestepBlock): + x = layer(x, emb) + elif isinstance(layer, AFNOCrossAttentionBlock3d): + img_shape = tuple(x.shape[-2:]) + x = layer(x, context[img_shape]) + else: + x = layer(x) + return x + + +class Upsample(nn.Module): + """ + An upsampling layer with an optional convolution. + :param channels: channels in the inputs and outputs. + :param use_conv: a bool determining if a convolution is applied. + :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then + upsampling occurs in the inner-two dimensions. + """ + + def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.dims = dims + if use_conv: + self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding) + + def forward(self, x): + assert x.shape[1] == self.channels + if self.dims == 3: + x = F.interpolate( + x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest" + ) + else: + x = F.interpolate(x, scale_factor=2, mode="nearest") + if self.use_conv: + x = self.conv(x) + return x + + +class Downsample(nn.Module): + """ + A downsampling layer with an optional convolution. + :param channels: channels in the inputs and outputs. + :param use_conv: a bool determining if a convolution is applied. + :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then + downsampling occurs in the inner-two dimensions. + """ + + def __init__(self, channels, use_conv, dims=2, out_channels=None,padding=1): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.dims = dims + stride = 2 if dims != 3 else (1, 2, 2) + if use_conv: + self.op = conv_nd( + dims, self.channels, self.out_channels, 3, stride=stride, padding=padding + ) + else: + assert self.channels == self.out_channels + self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride) + + def forward(self, x): + assert x.shape[1] == self.channels + return self.op(x) + + +class ResBlock(TimestepBlock): + """ + A residual block that can optionally change the number of channels. + :param channels: the number of input channels. + :param emb_channels: the number of timestep embedding channels. + :param dropout: the rate of dropout. + :param out_channels: if specified, the number of out channels. + :param use_conv: if True and out_channels is specified, use a spatial + convolution instead of a smaller 1x1 convolution to change the + channels in the skip connection. + :param dims: determines if the signal is 1D, 2D, or 3D. + :param use_checkpoint: if True, use gradient checkpointing on this module. + :param up: if True, use this block for upsampling. + :param down: if True, use this block for downsampling. + """ + + def __init__( + self, + channels, + emb_channels, + dropout, + out_channels=None, + use_conv=False, + use_scale_shift_norm=False, + dims=2, + use_checkpoint=False, + up=False, + down=False, + ): + super().__init__() + self.channels = channels + self.emb_channels = emb_channels + self.dropout = dropout + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_checkpoint = use_checkpoint + self.use_scale_shift_norm = use_scale_shift_norm + + self.in_layers = nn.Sequential( + normalization(channels), + nn.SiLU(), + conv_nd(dims, channels, self.out_channels, 3, padding=1), + ) + + self.updown = up or down + + if up: + self.h_upd = Upsample(channels, False, dims) + self.x_upd = Upsample(channels, False, dims) + elif down: + self.h_upd = Downsample(channels, False, dims) + self.x_upd = Downsample(channels, False, dims) + else: + self.h_upd = self.x_upd = nn.Identity() + + self.emb_layers = nn.Sequential( + nn.SiLU(), + linear( + emb_channels, + 2 * self.out_channels if use_scale_shift_norm else self.out_channels, + ), + ) + self.out_layers = nn.Sequential( + normalization(self.out_channels), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module( + conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1) + ), + ) + + if self.out_channels == channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = conv_nd( + dims, channels, self.out_channels, 3, padding=1 + ) + else: + self.skip_connection = conv_nd(dims, channels, self.out_channels, 1) + + def forward(self, x, emb): + """ + Apply the block to a Tensor, conditioned on a timestep embedding. + :param x: an [N x C x ...] Tensor of features. + :param emb: an [N x emb_channels] Tensor of timestep embeddings. + :return: an [N x C x ...] Tensor of outputs. + """ + return checkpoint( + self._forward, (x, emb), self.parameters(), self.use_checkpoint + ) + + + def _forward(self, x, emb): + if self.updown: + in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] + h = in_rest(x) + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + else: + h = self.in_layers(x) + emb_out = self.emb_layers(emb).type(h.dtype) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + if self.use_scale_shift_norm: + out_norm, out_rest = self.out_layers[0], self.out_layers[1:] + scale, shift = th.chunk(emb_out, 2, dim=1) + h = out_norm(h) * (1 + scale) + shift + h = out_rest(h) + else: + h = h + emb_out + h = self.out_layers(h) + return self.skip_connection(x) + h + + +class UNetModel(nn.Module): + """ + The full UNet model with attention and timestep embedding. + :param in_channels: channels in the input Tensor. + :param model_channels: base channel count for the model. + :param out_channels: channels in the output Tensor. + :param num_res_blocks: number of residual blocks per downsample. + :param attention_resolutions: a collection of downsample rates at which + attention will take place. May be a set, list, or tuple. + For example, if this contains 4, then at 4x downsampling, attention + will be used. + :param dropout: the dropout probability. + :param channel_mult: channel multiplier for each level of the UNet. + :param conv_resample: if True, use learned convolutions for upsampling and + downsampling. + :param dims: determines if the signal is 1D, 2D, or 3D. + :param use_checkpoint: use gradient checkpointing to reduce memory usage. + :param num_heads: the number of attention heads in each attention layer. + :param num_heads_channels: if specified, ignore num_heads and instead use + a fixed channel width per attention head. + :param num_heads_upsample: works with num_heads to set a different number + of heads for upsampling. Deprecated. + :param use_scale_shift_norm: use a FiLM-like conditioning mechanism. + :param resblock_updown: use residual blocks for up/downsampling. + + """ + + def __init__( + self, + model_channels, + in_channels=1, + out_channels=1, + num_res_blocks=2, + attention_resolutions=(1,2,4), + context_ch=128, + dropout=0, + channel_mult=(1, 2, 4, 4), + conv_resample=True, + dims=3, + use_checkpoint=False, + use_fp16=False, + num_heads=-1, + num_head_channels=-1, + num_heads_upsample=-1, + use_scale_shift_norm=False, + resblock_updown=False, + legacy=True, + num_timesteps=1 + ): + super().__init__() + + if num_heads_upsample == -1: + num_heads_upsample = num_heads + + if num_heads == -1: + assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set' + + if num_head_channels == -1: + assert num_heads != -1, 'Either num_heads or num_head_channels has to be set' + + self.in_channels = in_channels + self.model_channels = model_channels + self.out_channels = out_channels + self.num_res_blocks = num_res_blocks + self.attention_resolutions = attention_resolutions + self.dropout = dropout + self.channel_mult = channel_mult + self.conv_resample = conv_resample + self.use_checkpoint = use_checkpoint + self.dtype = th.float16 if use_fp16 else th.float32 + self.num_heads = num_heads + self.num_head_channels = num_head_channels + self.num_heads_upsample = num_heads_upsample + timesteps = th.arange(1, num_timesteps+1) + + time_embed_dim = model_channels * 4 + self.time_embed = nn.Sequential( + linear(model_channels, time_embed_dim), + nn.SiLU(), + linear(time_embed_dim, time_embed_dim), + ) + + self.input_blocks = nn.ModuleList( + [ + TimestepEmbedSequential( + conv_nd(dims, in_channels, model_channels, 3, padding=1) + ) + ] + ) + self._feature_size = model_channels + input_block_chans = [model_channels] + ch = model_channels + ds = 1 + for level, mult in enumerate(channel_mult): + for _ in range(num_res_blocks): + layers = [ + ResBlock( + ch, + time_embed_dim, + dropout, + out_channels=mult * model_channels, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + ) + ] + ch = mult * model_channels + if ds in attention_resolutions: + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + if legacy: + dim_head = num_head_channels + layers.append( + AFNOCrossAttentionBlock3d( + ch, context_dim=context_ch[level], num_blocks=num_heads, + data_format="channels_first", timesteps=timesteps + ) + ) + self.input_blocks.append(TimestepEmbedSequential(*layers)) + self._feature_size += ch + input_block_chans.append(ch) + if level != len(channel_mult) - 1: + out_ch = ch + self.input_blocks.append( + TimestepEmbedSequential( + ResBlock( + ch, + time_embed_dim, + dropout, + out_channels=out_ch, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + down=True, + ) + if resblock_updown + else Downsample( + ch, conv_resample, dims=dims, out_channels=out_ch + ) + ) + ) + ch = out_ch + input_block_chans.append(ch) + ds *= 2 + self._feature_size += ch + + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + if legacy: + dim_head = num_head_channels + self.middle_block = TimestepEmbedSequential( + ResBlock( + ch, + time_embed_dim, + dropout, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + ), + AFNOCrossAttentionBlock3d( + ch, context_dim=context_ch[-1], num_blocks=num_heads, + data_format="channels_first", timesteps=timesteps + ), + ResBlock( + ch, + time_embed_dim, + dropout, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + ), + ) + self._feature_size += ch + + self.output_blocks = nn.ModuleList([]) + for level, mult in list(enumerate(channel_mult))[::-1]: + for i in range(num_res_blocks + 1): + ich = input_block_chans.pop() + layers = [ + ResBlock( + ch + ich, + time_embed_dim, + dropout, + out_channels=model_channels * mult, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + ) + ] + ch = model_channels * mult + if ds in attention_resolutions: + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + if legacy: + #num_heads = 1 + dim_head = num_head_channels + layers.append( + AFNOCrossAttentionBlock3d( + ch, context_dim=context_ch[level], num_blocks=num_heads, + data_format="channels_first", timesteps=timesteps + ) + ) + if level and i == num_res_blocks: + out_ch = ch + layers.append( + ResBlock( + ch, + time_embed_dim, + dropout, + out_channels=out_ch, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + up=True, + ) + if resblock_updown + else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch) + ) + ds //= 2 + self.output_blocks.append(TimestepEmbedSequential(*layers)) + self._feature_size += ch + + self.out = nn.Sequential( + normalization(ch), + nn.SiLU(), + zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)), + ) + + def forward(self, x, timesteps=None, context=None): + """ + Apply the model to an input batch. + :param x: an [N x C x ...] Tensor of inputs. + :param timesteps: a 1-D batch of timesteps. + :param context: conditioning plugged in via crossattn + :return: an [N x C x ...] Tensor of outputs. + """ + hs = [] + t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False) + emb = self.time_embed(t_emb) + + h = x.type(self.dtype) + for module in self.input_blocks: + h = module(h, emb, context) + hs.append(h) + h = self.middle_block(h, emb, context) + for module in self.output_blocks: + h = th.cat([h, hs.pop()], dim=1) + h = module(h, emb, context) + h = h.type(x.dtype) + return self.out(h) + diff --git a/ldcast/models/nowcast/nowcast.py b/ldcast/models/nowcast/nowcast.py new file mode 100644 index 0000000000000000000000000000000000000000..d2949f43f488fdd1fcbd9ea748466324f1681baa --- /dev/null +++ b/ldcast/models/nowcast/nowcast.py @@ -0,0 +1,256 @@ +import collections + +import torch +from torch import nn +from torch.nn import functional as F +import pytorch_lightning as pl + +from ..blocks.afno import AFNOBlock3d +from ..blocks.attention import positional_encoding, TemporalTransformer + + +class Nowcaster(pl.LightningModule): + def __init__(self, nowcast_net): + super().__init__() + self.nowcast_net = nowcast_net + + def forward(self, x): + return self.nowcast_net(x) + + def _loss(self, batch): + (x,y) = batch + y_pred = self.forward(x) + return (y-y_pred).square().mean() + + def training_step(self, batch, batch_idx): + loss = self._loss(batch) + self.log("train_loss", loss) + return loss + + @torch.no_grad() + def val_test_step(self, batch, batch_idx, split="val"): + loss = self._loss(batch) + log_params = {"on_step": False, "on_epoch": True, "prog_bar": True} + self.log(f"{split}_loss", loss, **log_params) + + def validation_step(self, batch, batch_idx): + self.val_test_step(batch, batch_idx, split="val") + + def test_step(self, batch, batch_idx): + self.val_test_step(batch, batch_idx, split="test") + + def configure_optimizers(self): + optimizer = torch.optim.AdamW( + self.parameters(), lr=1e-3, + betas=(0.5, 0.9), weight_decay=1e-3 + ) + reduce_lr = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, patience=3, factor=0.25, verbose=True + ) + + optimizer_spec = { + "optimizer": optimizer, + "lr_scheduler": { + "scheduler": reduce_lr, + "monitor": "val_loss", + "frequency": 1, + }, + } + return optimizer_spec + + +class AFNONowcastNetBasic(nn.Sequential): + def __init__( + self, + embed_dim=256, + depth=12, + patch_size=(4,4,4) + ): + patch_embed = PatchEmbed3d( + embed_dim=embed_dim, patch_size=patch_size + ) + blocks = nn.Sequential( + *(AFNOBlock(embed_dim) for _ in range(depth)) + ) + patch_expand = PatchExpand3d( + embed_dim=embed_dim, patch_size=patch_size + ) + super().__init__(*[patch_embed, blocks, patch_expand]) + + +class FusionBlock3d(nn.Module): + def __init__(self, dim, size_ratios, dim_out=None, afno_fusion=False): + super().__init__() + + N_sources = len(size_ratios) + if not isinstance(dim, collections.abc.Sequence): + dim = (dim,) * N_sources + if dim_out is None: + dim_out = dim[0] + + self.scale = nn.ModuleList() + for (i,size_ratio) in enumerate(size_ratios): + if size_ratio == 1: + scale = nn.Identity() + else: + scale = [] + while size_ratio > 1: + scale.append(nn.ConvTranspose3d( + dim[i], dim_out if size_ratio==2 else dim[i], + kernel_size=(1,3,3), stride=(1,2,2), + padding=(0,1,1), output_padding=(0,1,1) + )) + size_ratio //= 2 + scale = nn.Sequential(*scale) + self.scale.append(scale) + + self.afno_fusion = afno_fusion + + if self.afno_fusion: + if N_sources > 1: + self.fusion = nn.Sequential( + nn.Linear(sum(dim), sum(dim)), + AFNOBlock3d(dim*N_sources, mlp_ratio=2), + nn.Linear(sum(dim), dim_out) + ) + else: + self.fusion = nn.Identity() + + def resize_proj(self, x, i): + x = x.permute(0,4,1,2,3) + x = self.scale[i](x) + x = x.permute(0,2,3,4,1) + return x + + def forward(self, x): + x = [self.resize_proj(xx, i) for (i, xx) in enumerate(x)] + if self.afno_fusion: + x = torch.concat(x, axis=-1) + x = self.fusion(x) + else: + x = sum(x) + return x + + +class AFNONowcastNetBase(nn.Module): + def __init__( + self, + autoencoder, + embed_dim=128, + embed_dim_out=None, + analysis_depth=4, + forecast_depth=4, + input_patches=(1,), + input_size_ratios=(1,), + output_patches=2, + train_autoenc=False, + afno_fusion=False + ): + super().__init__() + + self.train_autoenc = train_autoenc + if not isinstance(autoencoder, collections.abc.Sequence): + autoencoder = [autoencoder] + if not isinstance(input_patches, collections.abc.Sequence): + input_patches = [input_patches] + num_inputs = len(autoencoder) + if not isinstance(embed_dim, collections.abc.Sequence): + embed_dim = [embed_dim] * num_inputs + if embed_dim_out is None: + embed_dim_out = embed_dim[0] + if not isinstance(analysis_depth, collections.abc.Sequence): + analysis_depth = [analysis_depth] * num_inputs + self.embed_dim = embed_dim + self.embed_dim_out = embed_dim_out + self.output_patches = output_patches + + # encoding + analysis for each input + self.autoencoder = nn.ModuleList() + self.proj = nn.ModuleList() + self.analysis = nn.ModuleList() + for i in range(num_inputs): + ae = autoencoder[i].requires_grad_(train_autoenc) + self.autoencoder.append(ae) + + proj = nn.Conv3d(ae.hidden_width, embed_dim[i], kernel_size=1) + self.proj.append(proj) + + analysis = nn.Sequential( + *(AFNOBlock3d(embed_dim[i]) for _ in range(analysis_depth[i])) + ) + self.analysis.append(analysis) + + # temporal transformer + self.use_temporal_transformer = \ + any((ipp != output_patches) for ipp in input_patches) + if self.use_temporal_transformer: + self.temporal_transformer = nn.ModuleList( + TemporalTransformer(embed_dim[i]) for i in range(num_inputs) + ) + + # data fusion + self.fusion = FusionBlock3d(embed_dim, input_size_ratios, + afno_fusion=afno_fusion, dim_out=embed_dim_out) + + # forecast + self.forecast = nn.Sequential( + *(AFNOBlock3d(embed_dim_out) for _ in range(forecast_depth)) + ) + + def add_pos_enc(self, x, t): + if t.shape[1] != x.shape[1]: + # this can happen if x has been compressed + # by the autoencoder in the time dimension + ds_factor = t.shape[1] // x.shape[1] + t = F.avg_pool1d(t.unsqueeze(1), ds_factor)[:,0,:] + + pos_enc = positional_encoding(t, x.shape[-1], add_dims=(2,3)) + return x + pos_enc + + def forward(self, x): + (x, t_relative) = list(zip(*x)) + + # encoding + analysis for each input + def process_input(i): + z = self.autoencoder[i].encode(x[i])[0] + z = self.proj[i](z) + z = z.permute(0,2,3,4,1) + z = self.analysis[i](z) + if self.use_temporal_transformer: + # add positional encoding + z = self.add_pos_enc(z, t_relative[i]) + + # transform to output shape and coordinates + expand_shape = z.shape[:1] + (-1,) + z.shape[2:] + pos_enc_output = positional_encoding( + torch.arange(1,self.output_patches+1, device=z.device), + self.embed_dim[i], add_dims=(0,2,3) + ) + pe_out = pos_enc_output.expand(*expand_shape) + z = self.temporal_transformer[i](pe_out, z) + return z + + x = [process_input(i) for i in range(len(x))] + + # merge inputs + x = self.fusion(x) + # produce prediction + x = self.forecast(x) + return x.permute(0,4,1,2,3) # to channels-first order + + +class AFNONowcastNet(AFNONowcastNetBase): + def __init__(self, autoencoder, output_autoencoder=None, **kwargs): + super().__init__(autoencoder, **kwargs) + if output_autoencoder is None: + output_autoencoder = autoencoder[0] + self.output_autoencoder = output_autoencoder.requires_grad_( + self.train_autoenc) + self.out_proj = nn.Conv3d( + self.embed_dim_out, output_autoencoder.hidden_width, kernel_size=1 + ) + + def forward(self, x): + x = super().forward(x) + x = self.out_proj(x) + return self.output_autoencoder.decode(x) diff --git a/ldcast/models/utils.py b/ldcast/models/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7894f57c04db3fd6b6e2bbdadc492ac8769a575b --- /dev/null +++ b/ldcast/models/utils.py @@ -0,0 +1,28 @@ +import torch +from torch import nn + + +def normalization(channels, norm_type="group", num_groups=32): + if norm_type == "batch": + return nn.BatchNorm3d(channels) + elif norm_type == "group": + return nn.GroupNorm(num_groups=num_groups, num_channels=channels) + elif (not norm_type) or (norm_type.tolower() == 'none'): + return nn.Identity() + else: + raise NotImplementedError(norm) + + +def activation(act_type="swish"): + if act_type == "swish": + return nn.SiLU() + elif act_type == "gelu": + return nn.GELU() + elif act_type == "relu": + return nn.ReLU() + elif act_type == "tanh": + return nn.Tanh() + elif not act_type: + return nn.Identity() + else: + raise NotImplementedError(act_type) diff --git a/ldcast/visualization/cm.py b/ldcast/visualization/cm.py new file mode 100644 index 0000000000000000000000000000000000000000..31614e7c636d47cba3fcadf597ab68b9a0ae2ddc --- /dev/null +++ b/ldcast/visualization/cm.py @@ -0,0 +1,36 @@ +import numpy as np +from matplotlib.colors import LinearSegmentedColormap + + +def yuv_rainbow_24(nc): + """ From https://github.com/ARM-DOE/pyart/blob/main/pyart/graph/_cm_colorblind.py + """ + path1 = np.linspace(0.8*np.pi, 1.8*np.pi, nc) + path2 = np.linspace(-0.33*np.pi, 0.33*np.pi, nc) + + y = np.concatenate([np.linspace(0.3, 0.85, nc*2//5), + np.linspace(0.9, 0.0, nc - nc*2//5)]) + u = 0.40*np.sin(path1) + v = 0.55*np.sin(path2) + 0.1 + + rgb_from_yuv = np.array([[1, 0, 1.13983], + [1, -0.39465, -0.58060], + [1, 2.03211, 0]]) + cmap_dict = {'blue': [], 'green': [], 'red': []} + for i in range(len(y)): + yuv = np.array([y[i], u[i], v[i]]) + rgb = rgb_from_yuv.dot(yuv) + red_tuple = (i/(len(y)-1.0), rgb[0], rgb[0]) + green_tuple = (i/(len(y)-1.0), rgb[1], rgb[1]) + blue_tuple = (i/(len(y)-1.0), rgb[2], rgb[2]) + cmap_dict['blue'].append(blue_tuple) + cmap_dict['red'].append(red_tuple) + cmap_dict['green'].append(green_tuple) + + return cmap_dict + + +homeyer_rainbow = LinearSegmentedColormap( + "homeyer_rainbow", + yuv_rainbow_24(15) +) diff --git a/ldcast/visualization/plots.py b/ldcast/visualization/plots.py new file mode 100644 index 0000000000000000000000000000000000000000..cd759dbd1e900aa3eed8d85946a5ec875d72c3ad --- /dev/null +++ b/ldcast/visualization/plots.py @@ -0,0 +1,606 @@ +import concurrent +import multiprocessing +import netCDF4 +import os + +from matplotlib import gridspec, colors, pyplot as plt +import numpy as np +import torch + +from ..analysis import confmatrix, fss, rank +from ..features import io +from .cm import homeyer_rainbow + + +def reverse_transform_R(R, mean=-0.051, std=0.528): + return 10**(R*std + mean) + + +def plot_precip_image( + ax, R, + Rmin=0, Rmax=25, threshold_mmh=0.1, + transform_R=False, + grid_spacing=64 +): + if isinstance(R, torch.Tensor): + R = R.detach().numpy() + if transform_R: + R = reverse_transform_R(R, mean=mean, std=std) + Rmin = reverse_transform_R(Rmin) + Rmax = reverse_transform_R(Rmax) + if threshold_mmh: + Rmin = max(Rmin, threshold_mmh) + R[R < threshold_mmh] = np.nan + norm = colors.LogNorm(Rmin, Rmax) + ax.set_yticks(np.arange(0, R.shape[0], grid_spacing)) + ax.set_xticks(np.arange(0, R.shape[1], grid_spacing)) + ax.grid(which='major', alpha=0.35) + ax.tick_params(left=False, bottom=False, + labelleft=False, labelbottom=False) + + return ax.imshow(R, norm=norm, cmap=homeyer_rainbow) + + +def plot_autoencoder_reconstruction( + R, R_hat, samples=8, timesteps=4, + out_file=None +): + fig = plt.figure(figsize=(2*timesteps*2+0.5, samples*2), dpi=150) + + gs = gridspec.GridSpec( + samples, 2*timesteps+1, + width_ratios=(1,)*(2*timesteps)+(0.2,), + wspace=0.02, hspace=0.02 + ) + for k in range(samples): + for (i,j) in enumerate(range(-timesteps,0)): + ax = fig.add_subplot(gs[k,i]) + im = plot_precip_image(ax, R[k,0,j,:,:]) + for (i,j) in enumerate(range(-timesteps,0)): + ax = fig.add_subplot(gs[k,i+timesteps]) + im = plot_precip_image(ax, R_hat[k,0,j,:,:]) + + cax = fig.add_subplot(gs[:,-1]) + plt.colorbar(im, cax=cax) + + if out_file is not None: + out_file = fig.savefig(out_file, bbox_inches='tight') + plt.close(fig) + + +def plot_animation(x, y, out_dir, sample=0, fmt="{}_{:02d}.png"): + def plot_frame(R, label, timestep): + fig = plt.figure() + ax = fig.add_subplot() + im = plot_precip_image(ax, R[sample,0,timestep,:,:]) + fn = fmt.format(label, k) + fn = os.path.join(out_dir, fn) + fig.savefig(fn, bbox_inches='tight') + plt.close(fig) + + for k in range(x.shape[2]): + plot_frame(x, "x", k) + + for k in range(y.shape[2]): + plot_frame(y, "y", k) + + +model_colors = { + "mch-dgmr": "#E69F00", + "mch-pysteps": "#009E73", + "mch-iters=50-res=256": "#0072B2", + "mch-persistence": "#888888", + "dwd-dgmr": "#E69F00", + "dwd-pysteps": "#009E73", + "dwd-iters=50-res=256": "#0072B2", + "dwd-persistence": "#888888", + "pm-mch-dgmr": "#E69F00", + "pm-mch-pysteps": "#009E73", + "pm-mch-iters=50-res=256": "#0072B2", + "pm-dwd-dgmr": "#E69F00", + "pm-dwd-pysteps": "#009E73", + "pm-dwd-iters=50-res=256": "#0072B2", +} + +scale_linestyles = { + "1x1": "-", + "8x8": "--", + "64x64": ":" +} + + +def plot_crps( + log=False, + models=("iters=50-res=256", "dgmr", "pysteps"), + scales=("1x1", "8x8", "64x64"), + model_labels=("LDCast", "DGMR", "PySTEPS", "Persist."), + interval_mins=5, + out_fn=None, + ax=None, + add_xlabel=True, + add_ylabel=True, + add_legend=True, + crop_box=None +): + crps = {} + crps_name = "logcrps" if log else "crps" + for model in models: + crps[model] = {} + fn = f"../results/crps/{crps_name}-{model}.nc" + with netCDF4.Dataset(fn, 'r') as ds: + for scale in scales: + var = f"crps_pool{scale}" + crps_model_scale = np.array(ds[var][:], copy=False) + if crop_box is not None: + scale_int = int(scale.split("x")[0]) + crps_model_scale = crps_model_scale[ + ..., + crop_box[0][0]//scale_int:crop_box[0][1]//scale_int, + crop_box[1][0]//scale_int:crop_box[1][1]//scale_int + ] + crps[model][scale] = crps_model_scale.mean(axis=(0,1,3,4)) + del crps_model_scale + + if ax is None: + fig = plt.figure(figsize=(8,5)) + ax = fig.add_subplot() + + max_t = 0 + for (model, label) in zip(models, model_labels): + for scale in scales: + score = crps[model][scale] + color = model_colors[model] + linestyle = scale_linestyles[scale] + t = np.arange( + interval_mins, (len(score)+0.1)*interval_mins, interval_mins + ) + max_t = max(max_t, t[-1]) + ax.plot(t, score, color=color, linestyle=linestyle, + label=label) + + if add_legend: + plt.legend() + if add_xlabel: + plt.xlabel("Lead time [min]", fontsize=12) + if add_ylabel: + plt.ylabel( + "LogCRPS" if log else "CRPS [mm h$^\\mathrm{-1}$]", + fontsize=12 + ) + + ax.set_xlim((0, max_t)) + ylim = ax.get_ylim() + ylim = (0, ylim[1]) + ax.set_ylim(ylim) + ax.tick_params(axis='both', which='major', labelsize=12) + + if out_fn is not None: + fig.savefig(out_fn, bbox_inches='tight') + plt.close(fig) + + +def plot_rank_distribution( + models=("iters=50-res=256", "dgmr", "pysteps"), + scales=("1x1", "8x8", "64x64"), + model_labels=("LDCast", "DGMR", "PySTEPS"), + out_fn=None, + num_ensemble_members=32, + ax=None, + add_xlabel=True, + add_ylabel=True, + add_legend=True, + crop_box=None +): + rank_hist = {} + rank_KL = {} + for model in models: + fn = f"../results/ranks/ranks-{model}.nc" + with netCDF4.Dataset(fn, 'r') as ds: + for scale in scales: + var = f"ranks_pool{scale}" + ranks_model_scale = np.array(ds[var][:], copy=False) + if crop_box is not None: + scale_int = int(scale.split("x")[0]) + ranks_model_scale = ranks_model_scale[ + ..., + crop_box[0][0]//scale_int:crop_box[0][1]//scale_int, + crop_box[1][0]//scale_int:crop_box[1][1]//scale_int + ] + rank_hist[(model,scale)] = rank.rank_distribution(ranks_model_scale) + rank_KL[(model,scale)] = rank.rank_DKL(rank_hist[(model,scale)]) + del ranks_model_scale + + if ax is None: + fig = plt.figure(figsize=(8,5)) + ax = fig.add_subplot() + + for scale in scales: + linestyle = scale_linestyles[scale] + for (model, label) in zip(models, model_labels): + h = rank_hist[(model,scale)] + color = model_colors[model] + x = np.linspace(0, 1, num_ensemble_members+1) + label_with_score = f"{label}: {rank_KL[(model,scale)]:.3f}" + ax.plot(x, h, color=color, linestyle=linestyle, + label=label_with_score) + h_ideal = 1/(num_ensemble_members+1) + ax.plot([0, 1], [h_ideal, h_ideal], color=(0.4,0.4,0.4), + linewidth=1.0) + + if add_legend: + ax.legend(loc='upper center') + if add_xlabel: + ax.set_xlabel("Normalized rank", fontsize=12) + if add_ylabel: + ax.set_ylabel("Occurrence", fontsize=12) + + ax.set_xlim((0, 1)) + ylim = ax.get_ylim() + ylim = (0, ylim[1]) + ax.set_ylim(ylim) + ax.tick_params(axis='both', which='major', labelsize=12) + ax.set_xticks([0, 0.25, 0.5, 0.75, 1]) + # int labels for 0 and 1 to save space + ax.set_xticklabels(["0", "0.25", "0.5", "0.75", "1"]) + + if out_fn is not None: + fig.savefig(out_fn, bbox_inches='tight') + plt.close(fig) + + +def plot_rank_metric( + models=("iters=50-res=256", "dgmr", "pysteps"), + scales=("1x1", "8x8", "64x64"), + model_labels=("LDCast", "DGMR", "PySTEPS"), + interval_mins=5, + metric_name="KL", + out_fn=None, + ax=None, + add_xlabel=True, + add_ylabel=True, + add_legend=True, +): + rank_metric = {} + for model in models: + rank_metric[model] = {} + fn = f"../results/ranks/ranks-{model}.nc" + with netCDF4.Dataset(fn, 'r') as ds: + for scale in scales: + var = f"ranks_pool{scale}" + ranks = np.array(ds[var][:], copy=False) + rank_metric[model][scale] = rank.rank_metric_by_leadtime(ranks) + del ranks + + if ax is None: + fig = plt.figure(figsize=(8,5)) + ax = fig.add_subplot() + + max_t = 0 + for (model, label) in zip(models, model_labels): + for scale in scales: + score = rank_metric[model][scale] + color = model_colors[model] + linestyle = scale_linestyles[scale] + label_with_scale = f"{label} {scale}" + t = np.arange( + interval_mins, (len(score)+0.1)*interval_mins, interval_mins + ) + max_t = max(max_t, t[-1]) + ax.plot(t, score, color=color, linestyle=linestyle, + label=label_with_scale) + + if add_legend: + plt.legend() + if add_xlabel: + plt.xlabel("Lead time [min]", fontsize=12) + if add_ylabel: + plt.ylabel(f"Rank {metric_name}", fontsize=12) + + ax.set_xlim((0, max_t)) + ylim = ax.get_ylim() + ylim = (0, ylim[1]) + ax.set_ylim(ylim) + ax.tick_params(axis='both', which='major', labelsize=12) + + if out_fn is not None: + fig.savefig(out_fn, bbox_inches='tight') + plt.close(fig) + + +def load_fss(model, scale, use_timesteps, crop_box): + fn = f"../results/fractions/fractions-{model}.nc" + with netCDF4.Dataset(fn, 'r') as ds: + sn = f"{scale}x{scale}" + obs_frac = np.array(ds[f"obs_frac_scale{sn}"][:], copy=False) + fc_frac = np.array(ds[f"fc_frac_scale{sn}"][:], copy=False) + if crop_box is not None: + obs_frac = obs_frac[ + ..., + crop_box[0][0]//scale:crop_box[0][1]//scale, + crop_box[1][0]//scale:crop_box[1][1]//scale + ] + fc_frac = fc_frac[ + ..., + crop_box[0][0]//scale:crop_box[0][1]//scale, + crop_box[1][0]//scale:crop_box[1][1]//scale + ] + return fss.fractions_skill_score( + obs_frac, fc_frac, use_timesteps=use_timesteps + ) + + +def plot_fss( + log=False, + models=("iters=50-res=256", "dgmr", "pysteps"), + model_labels=("LDCast", "DGMR", "PySTEPS"), + interval_mins=5, + out_fn=None, + ax=None, + add_xlabel=True, + add_ylabel=True, + add_legend=True, + scales=None, + use_timesteps=18, + crop_box=None +): + if scales is None: + scales = 2**np.arange(9) + fss_scale = {} + N_threads = min(multiprocessing.cpu_count(), len(models)*len(scales)) + with concurrent.futures.ProcessPoolExecutor(N_threads) as executor: + for model in models: + fss_scale[model] = {} + for scale in scales: + fss_scale[model][scale] = executor.submit( + load_fss, model, scale, use_timesteps, crop_box + ) + + for model in models: + for scale in scales: + fss_scale[model][scale] = fss_scale[model][scale].result() + + if ax is None: + fig = plt.figure(figsize=(8,5)) + ax = fig.add_subplot() + + for (model, label) in zip(models, model_labels): + scales = sorted(fss_scale[model]) + fss_for_model = [fss_scale[model][s] for s in scales] + + model_parts = model.split("-") + if model.startswith("pm-"): + model_without_threshold = "-".join(model_parts[:1] + model_parts[2:]) + else: + model_without_threshold = "-".join(model_parts[1:]) + color = model_colors[model_without_threshold] + + ax.plot(scales, fss_for_model, color=color, + label=label) + + if add_legend: + plt.legend() + if add_xlabel: + plt.xlabel("Scale [km]", fontsize=12) + if add_ylabel: + plt.ylabel("FSS") + + ax.set_xlim((scales[0], scales[-1])) + ylim = ax.get_ylim() + ylim = (0, ylim[1]) + ax.set_ylim(ylim) + ax.tick_params(axis='both', which='major', labelsize=12) + + if out_fn is not None: + fig.savefig(out_fn, bbox_inches='tight') + plt.close(fig) + + +def plot_csi_threshold( + models=("iters=50-res=256", "dgmr", "pysteps"), + scales=("1x1", "8x8", "64x64"), + prob_thresholds=tuple(np.linspace(0,1,33)), + model_labels=("LDCast", "DGMR", "PySTEPS"), + out_fn=None, + num_ensemble_members=32, + max_timestep=18, + ax=None, + add_xlabel=True, + add_ylabel=True, + add_legend=True, + crop_box=None +): + csi = {} + for model in models: + fn = f"../results/fractions/fractions-{model}.nc" + with netCDF4.Dataset(fn, 'r') as ds: + for scale in scales: + fc_var = f"fc_frac_scale{scale}" + fc_frac = np.array(ds[fc_var], copy=False) + fc_frac = fc_frac[...,:max_timestep,:,:] + obs_var = f"obs_frac_scale{scale}" + obs_frac = np.array(ds[obs_var], copy=False) + obs_frac = obs_frac[...,:max_timestep,:,:] + conf_matrix = confmatrix.confusion_matrix_thresholds( + fc_frac, obs_frac, prob_thresholds + ) + del fc_frac, obs_frac + csi_scale = confmatrix.intersection_over_union(conf_matrix) + csi[(model,scale)] = csi_scale + + if ax is None: + fig = plt.figure(figsize=(8,5)) + ax = fig.add_subplot() + + for scale in scales: + linestyle = scale_linestyles[scale] + for (model, label) in zip(models, model_labels): + c = csi[(model,scale)] + model_parts = model.split("-") + if model.startswith("pm-"): + model_without_threshold = "-".join(model_parts[:1] + model_parts[2:]) + else: + model_without_threshold = "-".join(model_parts[1:]) + color = model_colors[model_without_threshold] + ax.plot(prob_thresholds, c, color=color, linestyle=linestyle, label=label) + + if add_legend: + ax.legend(loc='upper center') + if add_xlabel: + ax.set_xlabel("Prob. threshold", fontsize=12) + if add_ylabel: + ax.set_ylabel("CSI", fontsize=12) + + ax.set_xlim((0, 1)) + ylim = ax.get_ylim() + ylim = (0, ylim[1]) + ax.set_ylim(ylim) + ax.tick_params(axis='both', which='major', labelsize=12) + ax.set_xticks([0, 0.25, 0.5, 0.75, 1]) + # int labels for 0 and 1 to save space + ax.set_xticklabels(["0", "0.25", "0.5", "0.75", "1"]) + + if out_fn is not None: + fig.savefig(out_fn, bbox_inches='tight') + plt.close(fig) + + +def plot_csi_leadtime( + models=("iters=50-res=256", "dgmr", "pysteps"), + scales=("1x1", "8x8", "64x64"), + prob_thresholds=tuple(np.linspace(0,1,33)), + model_labels=("LDCast", "DGMR", "PySTEPS"), + out_fn=None, + interval_mins=5, + num_ensemble_members=32, + ax=None, + add_xlabel=True, + add_ylabel=True, + add_legend=True, + crop_box=None +): + csi = {} + for model in models: + fn = f"../results/fractions/fractions-{model}.nc" + with netCDF4.Dataset(fn, 'r') as ds: + for scale in scales: + fc_var = f"fc_frac_scale{scale}" + fc_frac = np.array(ds[fc_var], copy=False) + obs_var = f"obs_frac_scale{scale}" + obs_frac = np.array(ds[obs_var], copy=False) + conf_matrix = confmatrix.confusion_matrix_thresholds_leadtime( + fc_frac, obs_frac, prob_thresholds + ) + + csi_scale = confmatrix.intersection_over_union(conf_matrix) + csi[(model,scale)] = np.nanmax(csi_scale, axis=1) + + max_t = 0 + for (model, label) in zip(models, model_labels): + for scale in scales: + score = csi[(model,scale)] + + model_parts = model.split("-") + if model.startswith("pm-"): + model_without_threshold = "-".join(model_parts[:1] + model_parts[2:]) + else: + model_without_threshold = "-".join(model_parts[1:]) + color = model_colors[model_without_threshold] + linestyle = scale_linestyles[scale] + t = np.arange( + interval_mins, (len(score)+0.1)*interval_mins, interval_mins + ) + max_t = max(max_t, t[-1]) + ax.plot(t, score, color=color, linestyle=linestyle, + label=label) + + if add_legend: + plt.legend() + if add_xlabel: + plt.xlabel("Lead time [min]", fontsize=12) + if add_ylabel: + plt.ylabel("CSI", fontsize=12) + + ax.set_xlim((0, max_t)) + ylim = ax.get_ylim() + ylim = (0, ylim[1]) + ax.set_ylim(ylim) + ax.tick_params(axis='both', which='major', labelsize=12) + + if out_fn is not None: + fig.savefig(out_fn, bbox_inches='tight') + plt.close(fig) + + +def plot_cost_loss_value( + models=("iters=50-res=256", "dgmr", "pysteps"), + scales=("1x1", "8x8", "64x64"), + prob_thresholds=tuple(np.linspace(0,1,33)), + model_labels=("LDCast", "DGMR", "PySTEPS"), + out_fn=None, + interval_mins=5, + num_ensemble_members=32, + ax=None, + add_xlabel=True, + add_ylabel=True, + add_legend=True, + crop_box=None +): + value = {} + loss = 1.0 + cost = np.linspace(0.01, 1, 100) + for model in models: + fn = f"../results/fractions/fractions-{model}.nc" + with netCDF4.Dataset(fn, 'r') as ds: + for scale in scales: + fc_var = f"fc_frac_scale{scale}" + fc_frac = np.array(ds[fc_var], copy=False) + obs_var = f"obs_frac_scale{scale}" + obs_frac = np.array(ds[obs_var], copy=False) + conf_matrix = confmatrix.confusion_matrix_thresholds( + fc_frac, obs_frac, prob_thresholds + ) + + p_clim = obs_frac.mean() + value_scale = [] + for c in cost: + v = confmatrix.cost_loss_value( + conf_matrix, c, loss, p_clim + ) + value_scale.append(v[len(v)//2]) + value[(model,scale)] = np.array(value_scale) + + max_score = 0 + for (model, label) in zip(models, model_labels): + for scale in scales: + score = value[(model,scale)] + max_score = max(max_score, score[np.isfinite(score)].max()) + + model_parts = model.split("-") + if model.startswith("pm-"): + model_without_threshold = "-".join(model_parts[:1] + model_parts[2:]) + else: + model_without_threshold = "-".join(model_parts[1:]) + color = model_colors[model_without_threshold] + linestyle = scale_linestyles[scale] + + ax.plot(cost, score, color=color, linestyle=linestyle, + label=label) + + if add_legend: + plt.legend() + if add_xlabel: + plt.xlabel("Cost/loss ratio", fontsize=12) + if add_ylabel: + plt.ylabel("Value", fontsize=12) + + ax.set_xlim((0, 1)) + ylim = (0, max_score*1.05) + ax.set_ylim(ylim) + ax.tick_params(axis='both', which='major', labelsize=12) + ax.set_xticks([0, 0.25, 0.5, 0.75, 1]) + # int labels for 0 and 1 to save space + ax.set_xticklabels(["0", "0.25", "0.5", "0.75", "1"]) + + if out_fn is not None: + fig.savefig(out_fn, bbox_inches='tight') + plt.close(fig) diff --git a/models/.keep b/models/.keep new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/autoenc/autoenc-32-0.01.pt b/models/autoenc/autoenc-32-0.01.pt new file mode 100644 index 0000000000000000000000000000000000000000..3ffb2315591896ebbd8f8a3b8e465463815290ec --- /dev/null +++ b/models/autoenc/autoenc-32-0.01.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aa5ad4b8689aadbf702376e7afe5cb437ef5057675e78a8986837e8f28b3126e +size 1617490 diff --git a/models/autoenc/autoenc.pt b/models/autoenc/autoenc.pt new file mode 100644 index 0000000000000000000000000000000000000000..40b67936711c18f31f02949ef0693ceeb3358530 --- /dev/null +++ b/models/autoenc/autoenc.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bf3792a94ee4cf347ca498b10b851d58a0f2ce6b3062e7b59fec5761c7edbf24 +size 1616323 diff --git a/scripts/convert_data_NB_2nc.py b/scripts/convert_data_NB_2nc.py new file mode 100644 index 0000000000000000000000000000000000000000..c8ba41ae53f66b67602ff07bdc556a92b03cf2de --- /dev/null +++ b/scripts/convert_data_NB_2nc.py @@ -0,0 +1,76 @@ +import os +import numpy as np +import dask.array as da +import xarray as xr + +def load_all_file(data_dir=""): + data_list = [] + filtered_files = [] + for filename in os.listdir(data_dir): + if filename.startswith("202306"): + filtered_files.append(filename) + # if filename.endswith("00.npy"): + # filtered_files.append(filename) + sorted_files = sorted(filtered_files) + for item in sorted_files: + sub_dir = os.path.join(data_dir) + pathfile = sub_dir + "/" + item + file = np.load(pathfile) + data_list.extend([file]) + + lon = np.arange(103.5, 109.2, 0.00892) + lat = np.arange(8, 13.75, 0.00899) + + return data_list + +def preprocess_data(data_list, out_dir=""): + patches = [] + + # Define patch size + patch_size = 32 + # new_array = xr.DataArray(np.array(data_list[0]), dims=("dim_0", "dim_1")) + # Iterate over the array to extract patches + for k in range(len(data_list)): + for i in range(0, 640, patch_size): + for j in range(0, 640, patch_size): + patch = data_list[k][i:i+patch_size, j:j+patch_size] + patches.append(patch) + + print(len(patches)) + data_shape = len(patches) + patches_array = np.array(patches, dtype=np.uint8) + temp_array = np.array(np.random.rand(data_shape, 2), dtype=np.uint16) + temp_array2 = np.arange(256, dtype=np.float32) + temp_array3 = np.arange(data_shape, dtype=np.int64) + + data_da = da.from_array(patches_array, chunks=(data_shape,32,32)) # Adjust chunk size as needed for your data + data_da2 = da.from_array(temp_array, chunks=(data_shape, 2)) + data_da3 = da.from_array(temp_array3, chunks=(data_shape, )) + data_da4 = da.from_array(temp_array2, chunks=(256, )) + + # Create xarray DataArray with DaskArray as its backend + patches = xr.DataArray(data_da, dims=("dim_patch", "dim_heigh", "dim_width")) + patch_coords = xr.DataArray(data_da2, dims=("dim_patch1", "dim_coord")) + patch_times = xr.DataArray(data_da3, dims=("dim_patch2")) + zero_patch_coords = xr.DataArray(data_da2, dims=("dim_zero_patch", "dim_coord")) + zero_patch_times = xr.DataArray(data_da3, dims=("dim_zero_patch1")) + scale = xr.DataArray(data_da4, dims=("dim_scale")) + + ds = patches.to_dataset(name = 'patches') + ds['patch_coords'] = patch_coords + ds['patch_times'] = patch_times + ds['zero_patch_coords'] = zero_patch_coords + ds['zero_patch_times'] = zero_patch_times + ds['scale'] = scale + + ds.attrs["zero_value"] = 1 + out_dir = out_dir + "/" + "RZC" + os.makedirs(out_dir, exist_ok=True) + file_name = os.path.join(out_dir, "patches_RV_202306.nc") + ds.to_netcdf(file_name) + + return len(data_list) + + +list = load_all_file(data_dir="/data/data_WF/ldcast_precipitation/test") +print(preprocess_data(list, out_dir="/data/data_WF/ldcast_precipitation/preprocess_data_test")) diff --git a/scripts/convert_radar_to_numpy.py b/scripts/convert_radar_to_numpy.py new file mode 100644 index 0000000000000000000000000000000000000000..5fbb4fbf0ebac696bf93c89a2a61e09ed5bc9683 --- /dev/null +++ b/scripts/convert_radar_to_numpy.py @@ -0,0 +1,102 @@ +from datetime import datetime, timedelta +import glob +import os + +from fire import Fire +import h5py +from matplotlib import pyplot as plt +import numpy as np +import pyart +import pathlib +import cartopy +import wradlib as wrl +import dask.array as da +import pandas as pd +import xarray as xr +import re + +from ldcast import forecast +from ldcast.visualization import plots + + +MAP_PROJECTION = cartopy.crs.PlateCarree() +#2020-12-31T17:10 +def convert_to_datetime(filename=""): + year = int(filename[0:4]) + month = int(filename[5:7]) + day = int(filename[8:10]) + hour = int(filename[11:13]) + minute = int(filename[14:16]) + return datetime(year,month,day,hour,minute) + +def filter_files_by_format(directory, pattern): + filtered_files = [] + regex = re.compile(pattern) + for filename in os.listdir(directory): + if regex.match(filename): + #if filename.endswith("00"): + # filtered_files.append(filename) + if filename.endswith("0"): + filtered_files.append(filename) + return filtered_files + +def read_multiple_data( + data_dir, + path_file +): + path_file = os.path.join(data_dir, path_file) + radar = pyart.io.read(path_file) + processed_grid = pyart.map.grid_from_radars( + radar, + # grid_shape=(1, 1000, 1000), + grid_shape=(20, 640, 640), + grid_limits=((1000, 5000), (-300_000, 300_000), (-300_000, 300_000)), + ) + reflectivity_data = processed_grid.fields['reflectivity']['data'].data + cmax_data = np.array(np.max(reflectivity_data, axis=0)) + radar_Z = wrl.trafo.idecibel(cmax_data) + R = wrl.zr.z_to_r(radar_Z, a=200.0, b=1.6) + processed_grid.fields['reflectivity']['data'].data[:] = R + R = processed_grid.fields['reflectivity']['data'].data[0] + # extra_left, extra_right = 128, 128 + # extra_top, extra_bottom = 32, 32 + + # R = np.pad(R, ((extra_top, extra_bottom), (extra_left, extra_right)), + # mode='constant', constant_values=0) + + return R + + + +def demo( + out_dir="/data/data_WF/ldcast_precipitation/reflectivity/test_GT", + data_dir="/data/data_WF/NhaBe/2023", + t0=datetime(2023,6,25,0,0), + interval=timedelta(minutes=10), + past_timesteps=4, +): + temp = ['Pro-Raw (01-10) T04-2023', 'Pro-Raw (11-20) T04-2023', 'Pro-Raw (21-30) T04-2023', 'Pro-Raw (01-08)-T5-2023 NHBE', 'Pro-Raw (09-15)-T5-2023 NHBE', 'Pro-Raw (16-23)-T5-2023 NHBE', + 'Pro-Raw (24-31)-T5-2023 NHBE', 'Pro-Raw (01-08) T06-2023' , 'Pro-Raw (09-16) T06-2023', 'Pro-Raw (17-24) T06-2023', 'Pro-Raw (25-30) T06-2023'] + # for subdir in os.listdir(directory_path): + for subdir in temp: + # for subdir in os.listdir(data_dir): + subdir_path = os.path.join(data_dir, subdir) + if os.path.isdir(subdir_path): + temp2 = os.listdir(subdir_path) + temp2.sort() + for subsubdir in temp2: + subsubdir_path = os.path.join(subdir_path, subsubdir) + pattern_to_filter = r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}" + filtered_files = filter_files_by_format(subsubdir_path, pattern_to_filter) + sorted_files = sorted(filtered_files) + print("Filtered Files:", len(filtered_files)) + for i in range(0, len(sorted_files)): + t = convert_to_datetime(sorted_files[i]) + print(t) + R_gt = read_multiple_data(subsubdir_path, sorted_files[i]) + fn = os.path.join(out_dir, f"{t:%Y%m%d%H%M}.npy") + np.save(fn, R_gt) + + +if __name__ == "__main__": + Fire(demo) diff --git a/scripts/convert_radar_to_numpy2.py b/scripts/convert_radar_to_numpy2.py new file mode 100644 index 0000000000000000000000000000000000000000..18325c1cc7d4699475b964af70228adaf91613be --- /dev/null +++ b/scripts/convert_radar_to_numpy2.py @@ -0,0 +1,102 @@ +from datetime import datetime, timedelta +import glob +import os + +from fire import Fire +import h5py +from matplotlib import pyplot as plt +import numpy as np +import pyart +import pathlib +import cartopy +import wradlib as wrl +import dask.array as da +import pandas as pd +import xarray as xr +import re + +from ldcast import forecast +from ldcast.visualization import plots + + +MAP_PROJECTION = cartopy.crs.PlateCarree() +#2020-12-31T17:10 +def convert_to_datetime(filename=""): + year = int(filename[0:4]) + month = int(filename[5:7]) + day = int(filename[8:10]) + hour = int(filename[11:13]) + minute = int(filename[14:16]) + return datetime(year,month,day,hour,minute) + +def filter_files_by_format(directory, pattern): + filtered_files = [] + regex = re.compile(pattern) + for filename in os.listdir(directory): + if regex.match(filename): + #if filename.endswith("00"): + # filtered_files.append(filename) + if filename.endswith("0"): + filtered_files.append(filename) + return filtered_files + +def read_multiple_data( + data_dir, + path_file +): + path_file = os.path.join(data_dir, path_file) + radar = pyart.io.read(path_file) + processed_grid = pyart.map.grid_from_radars( + radar, + # grid_shape=(1, 1000, 1000), + grid_shape=(20, 640, 640), + grid_limits=((1000, 5000), (-300_000, 300_000), (-300_000, 300_000)), + ) + reflectivity_data = processed_grid.fields['reflectivity']['data'].data + cmax_data = np.array(np.max(reflectivity_data, axis=0)) + radar_Z = wrl.trafo.idecibel(cmax_data) + R = wrl.zr.z_to_r(radar_Z, a=200.0, b=1.6) + processed_grid.fields['reflectivity']['data'].data[:] = R + R = processed_grid.fields['reflectivity']['data'].data[0] + # extra_left, extra_right = 128, 128 + # extra_top, extra_bottom = 32, 32 + + # R = np.pad(R, ((extra_top, extra_bottom), (extra_left, extra_right)), + # mode='constant', constant_values=0) + + return R + + + +def demo( + out_dir="/data/data_WF/ldcast_precipitation/reflectivity/train_GT", + data_dir="/data/data_WF/NhaBe/2022", + t0=datetime(2023,6,25,0,0), + interval=timedelta(minutes=10), + past_timesteps=4, +): + temp = [ 'T02-2022Pro-Raw (01-19)', 'T02-2022Pro-Raw (20-28)', 'T03-2022Pro-Raw (20-31)', 'T04-2022Pro-Raw (01-16)-', +'T04-2022Pro-Raw (17-30)-', 'T05-2022Pro-Raw (01-11)', 'T05-2022Pro-Raw (12-22)', 'T05-2022Pro-Raw (23-31)' , 'T06- 2022Pro-Raw (01-13)', 'T06- 2022Pro-Raw (14-23)', 'T06- 2022Pro-Raw (24-30)'] + # for subdir in os.listdir(directory_path): + for subdir in temp: + # for subdir in os.listdir(data_dir): + subdir_path = os.path.join(data_dir, subdir) + if os.path.isdir(subdir_path): + temp2 = os.listdir(subdir_path) + temp2.sort() + for subsubdir in temp2: + subsubdir_path = os.path.join(subdir_path, subsubdir) + pattern_to_filter = r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}" + filtered_files = filter_files_by_format(subsubdir_path, pattern_to_filter) + sorted_files = sorted(filtered_files) + print("Filtered Files:", len(filtered_files)) + for i in range(0, len(sorted_files)): + t = convert_to_datetime(sorted_files[i]) + print(t) + R_gt = read_multiple_data(subsubdir_path, sorted_files[i]) + fn = os.path.join(out_dir, f"{t:%Y%m%d%H%M}.npy") + np.save(fn, R_gt) + + +if __name__ == "__main__": + Fire(demo) diff --git a/scripts/convert_radar_to_numpy3.py b/scripts/convert_radar_to_numpy3.py new file mode 100644 index 0000000000000000000000000000000000000000..2f75fc96f7001fdaf13de18926a19ef1d84865b4 --- /dev/null +++ b/scripts/convert_radar_to_numpy3.py @@ -0,0 +1,103 @@ +from datetime import datetime, timedelta +import glob +import os + +from fire import Fire +import h5py +from matplotlib import pyplot as plt +import numpy as np +import pyart +import pathlib +import cartopy +import wradlib as wrl +import dask.array as da +import pandas as pd +import xarray as xr +import re + +from ldcast import forecast +from ldcast.visualization import plots + + +MAP_PROJECTION = cartopy.crs.PlateCarree() +#2020-12-31T17:10 +def convert_to_datetime(filename=""): + year = int(filename[0:4]) + month = int(filename[5:7]) + day = int(filename[8:10]) + hour = int(filename[11:13]) + minute = int(filename[14:16]) + return datetime(year,month,day,hour,minute) + +def filter_files_by_format(directory, pattern): + filtered_files = [] + regex = re.compile(pattern) + for filename in os.listdir(directory): + if regex.match(filename): + #if filename.endswith("00"): + # filtered_files.append(filename) + if filename.endswith("0"): + filtered_files.append(filename) + return filtered_files + +def read_multiple_data( + data_dir, + path_file +): + path_file = os.path.join(data_dir, path_file) + radar = pyart.io.read(path_file) + processed_grid = pyart.map.grid_from_radars( + radar, + # grid_shape=(1, 1000, 1000), + grid_shape=(20, 640, 640), + grid_limits=((1000, 5000), (-300_000, 300_000), (-300_000, 300_000)), + ) + reflectivity_data = processed_grid.fields['reflectivity']['data'].data + cmax_data = np.array(np.max(reflectivity_data, axis=0)) + radar_Z = wrl.trafo.idecibel(cmax_data) + R = wrl.zr.z_to_r(radar_Z, a=200.0, b=1.6) + processed_grid.fields['reflectivity']['data'].data[:] = R + R = processed_grid.fields['reflectivity']['data'].data[0] + # extra_left, extra_right = 128, 128 + # extra_top, extra_bottom = 32, 32 + + # R = np.pad(R, ((extra_top, extra_bottom), (extra_left, extra_right)), + # mode='constant', constant_values=0) + + return R + + + +def demo( + out_dir="/data/data_WF/ldcast_precipitation/reflectivity/train_GT", + data_dir="/data/data_WF/NhaBe/2022", + t0=datetime(2023,6,25,0,0), + interval=timedelta(minutes=10), + past_timesteps=4, +): + temp = ['T08-2022Pro-Raw (01-08)', 'T08-2022Pro-Raw (08-16)' , 'T08-2022Pro-Raw (16-31)', +'T08-2022Pro-Raw (17-31)', 'T09-2022Pro- Raw (01-19)', 'T09-2022Pro- Raw (20-25)', 'T09-2022Pro- Raw (26-30)', 'T10-2022Pro-Raw (01-08)', 'T10-2022Pro-Raw (09-16)', 'T10-2022Pro-Raw (17-24)', +'T10-2022Pro-Raw (25-31)', 'T11-2022Pro-Raw (01-10)', 'T11-2022Pro-Raw (11-20)', 'T11-2022Pro-Raw (20-30)', 'T12-2022Pro-Raw (01-10)', 'T12-2022Pro-Raw (11-20)', 'T12-2022Pro-Raw (21-31)'] + # for subdir in os.listdir(directory_path): + for subdir in temp: + # for subdir in os.listdir(data_dir): + subdir_path = os.path.join(data_dir, subdir) + if os.path.isdir(subdir_path): + temp2 = os.listdir(subdir_path) + temp2.sort() + for subsubdir in temp2: + subsubdir_path = os.path.join(subdir_path, subsubdir) + pattern_to_filter = r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}" + filtered_files = filter_files_by_format(subsubdir_path, pattern_to_filter) + sorted_files = sorted(filtered_files) + print("Filtered Files:", len(filtered_files)) + for i in range(0, len(sorted_files)): + t = convert_to_datetime(sorted_files[i]) + print(t) + R_gt = read_multiple_data(subsubdir_path, sorted_files[i]) + fn = os.path.join(out_dir, f"{t:%Y%m%d%H%M}.npy") + np.save(fn, R_gt) + + +if __name__ == "__main__": + Fire(demo) diff --git a/scripts/evaluate_res.py b/scripts/evaluate_res.py new file mode 100644 index 0000000000000000000000000000000000000000..cd8eab96e44adc9e0f02b7c296a3aa4f177e2dec --- /dev/null +++ b/scripts/evaluate_res.py @@ -0,0 +1,113 @@ +import os +import numpy as np + +def anomaly_correlation_coefficient(ground_truth, predictions): + # Tính toán mean của ground truth và predictions + num_variables = ground_truth.shape[1] + acc_values = [] + for i in range(num_variables): + observed_var = ground_truth[:, i] + predicted_var = predictions[:, i] + mean_observed = np.mean(observed_var) + mean_predicted = np.mean(predicted_var) + covariance = np.mean((observed_var - mean_observed) * (predicted_var - mean_predicted)) + std_observed = np.std(observed_var) + std_predicted = np.std(predicted_var) + acc = covariance / (std_observed * std_predicted) + acc_values.append(acc) + mean_acc = np.mean(acc_values) + return mean_acc + +def root_mean_square_error(ground_truth, predictions): + mse = np.mean((ground_truth - predictions)**2) + rmse = np.sqrt(mse) + return rmse + +def weighted_mean_square_error(ground_truth, predictions, weights): + mse = np.mean(weights * (ground_truth - predictions)**2) + return mse + +def root_weighted_mean_square_error(ground_truth, predictions, weights): + wmse = weighted_mean_square_error(ground_truth, predictions, weights) + rwmse = np.sqrt(wmse) + return rwmse + +def crps(predicted_probs, observed_value): + # Sắp xếp các xác suất dự đoán theo thứ tự tăng dần + predicted_probs_sorted = np.sort(predicted_probs) + + # Tính toán số lượng điểm xác suất + n = len(predicted_probs) + + # Tính toán CRPS bằng cách sử dụng công thức CRPS diskret + crps = 0 + for i in range(n): + crps += (predicted_probs_sorted[i] - (i + 0.5) / n)**2 + + crps *= 1 / n + + # Thêm thành phần phân phối Gaussian + crps += np.mean((predicted_probs - observed_value)**2) + + return crps + + +# Đường dẫn đến thư mục chứa file ground truth và predictions +#folder_ground_truth = '/data/data_WF/finetune/output/exp_GT' +#folder_predictions = '/data/data_WF/finetune/output/exp_pred' +folder_ground_truth = '/data/data_WF/finetune/output/x_GT' +folder_predictions = '/data/data_WF/finetune/output/x_pred' +#folder_ground_truth = '/data/data_WF/ldcast_precipitation/train_ground_truth' +#folder_predictions = '/data/data_WF/ldcast_precipitation/train' +folder_weights = '' + +# Danh sách các tên file trong thư mục ground truth và predictions +gt_files = os.listdir(folder_ground_truth) +pred_files = os.listdir(folder_predictions) +# weight_files = os.listdir(folder_weights) + +# Khởi tạo một danh sách để lưu trữ các giá trị MSE +mse_list = [] +acc_list = [] +rmse_list = [] +# wmse_list = [] +# rwmse_list = [] +crps_list = [] + +# Vòng lặp qua các file trong thư mục +for gt_file, pred_file in zip(gt_files, pred_files): + # Đọc dữ liệu từ file ground truth và file dự đoán + ground_truth = np.load(os.path.join(folder_ground_truth, gt_file)) + predictions = np.load(os.path.join(folder_predictions, pred_file)) + ground_truth = ground_truth * 1000 + predictions = predictions * 1000 + + mse = np.mean((ground_truth - predictions)**2) + acc = anomaly_correlation_coefficient(ground_truth, predictions) + rmse = root_mean_square_error(ground_truth, predictions) + crps_score = crps(predictions, ground_truth) + + # wmse = weighted_mean_square_error(ground_truth, predictions, weights) + # rwmse = root_weighted_mean_square_error(ground_truth, predictions, weights) + + # Thêm MSE vào danh sách + mse_list.append(mse) + acc_list.append(acc) + rmse_list.append(rmse) + # wmse_list.append(wmse) + # rwmse_list.append(rwmse) + crps_list.append(crps_score) +# Tính toán trung bình của các giá trị MSE +average_mse = np.mean(mse_list) +average_acc = np.mean(acc_list) +average_rmse = np.mean(rmse_list) +# average_wmse = np.mean(wmse_list) +# average_rwmse = np.mean(rwmse_list) +average_crps = np.mean(crps_list) + +print("Average Mean Square Error (MSE) across all files:", average_mse) +print("Average Anomaly Correlation Coefficient (ACC) across all files:", average_acc) +print("Average Root Mean Square Error (RMSE) across all files:", average_rmse) +# print("Average Weighted Mean Square Error (WMSE) across all files:", average_wmse) +# print("Average Root Weighted Mean Square Error (RWMSE) across all files:", average_rwmse) +print("Average Continuous Ranked Probability Score (CRPS) across all files:", average_crps) diff --git a/scripts/fintune_autoenc.py b/scripts/fintune_autoenc.py new file mode 100644 index 0000000000000000000000000000000000000000..feab41c0eea12f7739b73203fdeae85ac7026a89 --- /dev/null +++ b/scripts/fintune_autoenc.py @@ -0,0 +1,114 @@ +import gc +import gzip +import os +import pickle + +from fire import Fire +import numpy as np +from omegaconf import OmegaConf + +from ldcast.features import batch, patches, split, transform +from ldcast.models.autoenc import encoder, training + +#file_dir = os.path.dirname(os.path.abspath(__file__)) +file_dir = os.path.dirname("/data/data_WF/ldcast_precipitation/preprocess_data/") + +def setup_data( + var="RZC", + #var="01", + batch_size=32, + sampler_file=None, + num_timesteps=8, + chunks_file="../data/split_chunks.pkl.gz" +): + variables = { + var: { + "sources": [var], + "timesteps": np.arange(num_timesteps), + } + } + predictors = [var] # autoencoder: predictors == targets + target = var + raw_vars = [var] + raw = { + var: patches.load_all_patches( + #os.path.join(file_dir, f"../data/{var}/"), var + os.path.join(file_dir, f"{var}/"), var + ) + for var in raw_vars + } + #print(len(raw)) + + # Load pregenerated train/valid/test split data. + # These can be generated with features.split.get_chunks() + with gzip.open(os.path.join(file_dir, chunks_file), 'rb') as f: + chunks = pickle.load(f) + (raw, _) = split.train_valid_test_split(raw, var, chunks=chunks) + + variables[var]["transform"] = transform.default_rainrate_transform( + raw["train"][var]["scale"] + ) + + if sampler_file is None: + sampler_file = { + "train": "../cache/sampler_autoenc_train.pkl", + "valid": "../cache/sampler_autoenc_valid.pkl", + "test": "../cache/sampler_autoenc_test.pkl", + } + bins = np.exp(np.linspace(np.log(0.2), np.log(50), 10)) + datamodule = split.DataModule( + variables, raw, predictors, target, var, + sampling_bins=bins, batch_size=batch_size, + sampler_file=sampler_file, + valid_seed=1234, test_seed=2345 + ) + + gc.collect() + return datamodule + + +def setup_model( + model_dir=None +): + enc = encoder.SimpleConvEncoder() + dec = encoder.SimpleConvDecoder() + (autoencoder, trainer) = training.setup_autoenc_training( + encoder=enc, + decoder=dec, + model_dir=model_dir + ) + gc.collect() + return (autoencoder, trainer) + + +def train( + var="RZC", + #var="01", + batch_size=8, + sampler_file=None, + num_timesteps=8, + chunks_file="split_chunks.pkl.gz", + model_dir=None, + ckpt_path=None +): + print("Loading data...") + datamodule = setup_data( + var=var, batch_size=batch_size, sampler_file=sampler_file, + num_timesteps=num_timesteps, chunks_file=chunks_file + ) + + print("Setting up model...") + (model, trainer) = setup_model(model_dir=model_dir) + + print("Starting training...") + trainer.fit(model, datamodule=datamodule, ckpt_path=ckpt_path) + trainer.test(model, datamodule=datamodule) + +def main(config=None, **kwargs): + config = OmegaConf.load(config) if (config is not None) else {} + config.update(kwargs) + train(**config) + + +if __name__ == "__main__": + Fire(main) diff --git a/scripts/forecast_demo.py b/scripts/forecast_demo.py new file mode 100644 index 0000000000000000000000000000000000000000..1e845e42c7caff94ca4b9f15c0774f6d1caff04e --- /dev/null +++ b/scripts/forecast_demo.py @@ -0,0 +1,132 @@ +from datetime import datetime, timedelta +import glob +import os + +from fire import Fire +import h5py +from matplotlib import pyplot as plt +import numpy as np + +from ldcast import forecast +from ldcast.visualization import plots + + +def read_data( + data_dir="../data/demo/20210622", + t0=datetime(2021,6,22,18,35), + interval=timedelta(minutes=5), + past_timesteps=4, + crop_box=((128,480), (160,608)) +): + cb = crop_box + R_past = [] + t = t0 - (past_timesteps-1) * interval + for i in range(past_timesteps): + timestamp = t.strftime("%y%j%H%M") + fn = f"RZC{timestamp}VL.[80]01.h5" + fn = os.path.join(data_dir, fn) + found_files = glob.glob(fn) + if found_files: + fn = found_files[0] + else: + raise FileNotFoundError(f"Unable to find data file {fn}.") + with h5py.File(fn, 'r') as f: + R = f["dataset1"]["data1"]["data"][:] + R = R[cb[0][0]:cb[0][1], cb[1][0]:cb[1][1]] + R_past.append(R) + t += interval + + R_past = np.stack(R_past, axis=0) + return R_past + + +def plot_border(ax, crop_box=((128,480), (160,608))): + import shapefile + border = shapefile.Reader("../data/Border_CH.shp") + shapes = list(border.shapeRecords()) + for shape in shapes: + x = np.array([i[0]/1000. for i in shape.shape.points[:]]) + y = np.array([i[1]/1000. for i in shape.shape.points[:]]) + ax.plot( + x-crop_box[1][0]-255, 480-y-crop_box[0][0], + 'k', linewidth=1.0 + ) + + +def plot_frame(R, fn, draw_border=True, t=None, label=None): + fig = plt.figure(dpi=150) + ax = fig.add_subplot() + plots.plot_precip_image(ax, R) + if draw_border: + plot_border(ax) + if t is not None: + timestamp = "%Y-%m-%d %H:%M UTC" + if label is not None: + timestamp += f" ({label})" + ax.text( + 0.02, 0.98, t.strftime(timestamp), + horizontalalignment='left', verticalalignment='top', + transform=ax.transAxes + ) + + fig.savefig(fn, bbox_inches='tight') + plt.close(fig) + + +def forecast_demo( + ldm_weights_fn="../models/genforecast/genforecast-radaronly-256x256-20step.pt", + autoenc_weights_fn="../models/autoenc/autoenc-32-0.01.pt", + num_diffusion_iters=50, + out_dir="../figures/demo/", + data_dir="../data/demo/20210622", + t0=datetime(2021,6,22,18,35), + interval=timedelta(minutes=5), + past_timesteps=4, + crop_box=((128,480), (160,608)), + draw_border=True, + ensemble_members=1, +): + R_past = read_data( + data_dir=data_dir, t0=t0, interval=interval, + past_timesteps=past_timesteps, crop_box=crop_box + ) + if ensemble_members == 1: + fc = forecast.Forecast( + ldm_weights_fn=ldm_weights_fn, + autoenc_weights_fn=autoenc_weights_fn + ) + R_pred = fc( + R_past, + num_diffusion_iters=num_diffusion_iters + ) + elif ensemble_members > 1: + fc = forecast.ForecastDistributed( + ldm_weights_fn=ldm_weights_fn, + autoenc_weights_fn=autoenc_weights_fn, + ) + R_past = R_past.reshape((1,) + R_past.shape) + R_pred = fc( + R_past, + num_diffusion_iters=num_diffusion_iters, + ensemble_members=ensemble_members + ) + R_past = R_past[0,...] + R_pred = R_pred[0,...].mean(axis=-1) # compute ensemble mean + else: + raise ValueError("ensemble_members must be > 0") + + os.makedirs(out_dir, exist_ok=True) + for k in range(R_past.shape[0]): + fn = os.path.join(out_dir, f"R_past-{k:02d}.png") + t = t0 - (R_past.shape[0]-k-1) * interval + plot_frame(R_past[k,:,:], fn, draw_border=draw_border, + t=t, label="Real") + for k in range(R_pred.shape[0]): + fn = os.path.join(out_dir, f"R_pred-{k:02d}.png") + t = t0 + (k+1)*interval + plot_frame(R_pred[k,:,:], fn, draw_border=draw_border, + t=t, label="Predicted") + + +if __name__ == "__main__": + Fire(forecast_demo) diff --git a/scripts/genpredict_0.py b/scripts/genpredict_0.py new file mode 100644 index 0000000000000000000000000000000000000000..6b48661027fb5decd0c12af5c52726c57a6af29c --- /dev/null +++ b/scripts/genpredict_0.py @@ -0,0 +1,80 @@ +from datetime import datetime, timedelta +import glob +import os + +from fire import Fire +import h5py +from matplotlib import pyplot as plt +import numpy as np +import pathlib +import cartopy +import wradlib as wrl +import dask.array as da +import pandas as pd +import xarray as xr +import re +from datetime import datetime, timedelta +from ldcast import forecast +from ldcast.visualization import plots +import torch +MAP_PROJECTION = cartopy.crs.PlateCarree() + + +def get_date_time(name): + year=int(name[0:4]) + month=int(name[4:6]) + day=int(name[6:8]) + hour=int(name[8:10]) + minutes=int(name[10:12]) + return datetime(year,month,day,hour,minutes) + +def demo( + ldm_weights_fn="/data/data_WF/ablation/ablation_time/genforecast-radaronly-256x256-20step.pt", + autoenc_weights_fn="/home/mmhk20/weather_forecast/ldcast/models/autoenc/autoenc-32-0.01.pt", + num_diffusion_iters=50, + out_dir="/data/weather2025/temp_train/train/pred_rad", + data_dir="/data/weather2025/temp_train/train/rad", + t0=datetime(2020,12,31,21,40), + interval=timedelta(minutes=10), + past_timesteps=4, + crop_box=((0,640), (0,640)), + draw_border=False, + ensemble_members=1, +): + filtered_files = [] + for filename in os.listdir(data_dir): + if filename.endswith("00.npz"): + filtered_files.append(filename) + if filename.endswith("03.npz"): + filtered_files.append(filename) + sorted_files = sorted(filtered_files) + fc1 = forecast.Forecast( + ldm_weights_fn=ldm_weights_fn, + autoenc_weights_fn=autoenc_weights_fn, + future_timesteps=20, + gpu=2 + ) + index = 170 + for i in range(index, len(sorted_files)): + print(i) + print(sorted_files[i]) + name_datetime = get_date_time(sorted_files[i]) + temp_radar = [] + for i in range(0,4): + temp_time = name_datetime + timedelta(minutes = i*10 - 30) + temp_path = os.path.join(data_dir,temp_time.strftime('%Y%m%d%H%M')+'.npz') + if(os.path.exists(temp_path)): + temp_radar.append(temp_path) + if(len(temp_radar) == 4): + with torch.no_grad(): + R_past = torch.stack([torch.tensor(np.load(temp_radar[i])['precipitation']) for i in range(0, len(temp_radar))],dim=0) + R_past.to('cuda:2') + R_pred = fc1(R_past,num_diffusion_iters=num_diffusion_iters) + npz_data = {'precipitation':R_pred[-3]} + pred_time = name_datetime + timedelta(minutes = 180) + path_save = os.path.join(out_dir,pred_time.strftime('%Y%m%d%H%M')+'.npz') + print(path_save) + np.savez(path_save,**npz_data) +if __name__ == "__main__": + Fire(demo) +#31779 diff --git a/scripts/genpredict_1.py b/scripts/genpredict_1.py new file mode 100644 index 0000000000000000000000000000000000000000..3186944926968240ad7c9df1567db9fe08f53dc2 --- /dev/null +++ b/scripts/genpredict_1.py @@ -0,0 +1,81 @@ +from datetime import datetime, timedelta +import glob +import os + +from fire import Fire +import h5py +from matplotlib import pyplot as plt +import numpy as np +import pathlib +import cartopy +import wradlib as wrl +import dask.array as da +import pandas as pd +import xarray as xr +import re +from datetime import datetime, timedelta +from ldcast import forecast +from ldcast.visualization import plots +import torch +MAP_PROJECTION = cartopy.crs.PlateCarree() + + +def get_date_time(name): + year=int(name[0:4]) + month=int(name[4:6]) + day=int(name[6:8]) + hour=int(name[8:10]) + minutes=int(name[10:12]) + return datetime(year,month,day,hour,minutes) + +def demo( + ldm_weights_fn="/data/data_WF/ablation/ablation_time/genforecast-radaronly-256x256-20step.pt", + autoenc_weights_fn="/home/mmhk20/weather_forecast/ldcast/models/autoenc/autoenc-32-0.01.pt", + num_diffusion_iters=50, + out_dir="/data/AI102024/val/pred_rad", + data_dir="/data/AI102024/val/rad", + t0=datetime(2020,12,31,21,40), + interval=timedelta(minutes=10), + past_timesteps=4, + crop_box=((0,640), (0,640)), + draw_border=False, + ensemble_members=1, +): + filtered_files = [] + for filename in os.listdir(data_dir): + if filename.endswith("00.npz"): + filtered_files.append(filename) + if filename.endswith("03.npz"): + filtered_files.append(filename) + sorted_files = sorted(filtered_files) + fc1 = forecast.Forecast( + ldm_weights_fn=ldm_weights_fn, + autoenc_weights_fn=autoenc_weights_fn, + future_timesteps=20, + gpu=1 + ) + index = 0 + print(len(sorted_files)) + for i in range(index, len(sorted_files)): + print(i) + print(sorted_files[i]) + name_datetime = get_date_time(sorted_files[i]) + temp_radar = [] + for i in range(0,4): + temp_time = name_datetime + timedelta(minutes = i*10 - 30) + temp_path = os.path.join(data_dir,temp_time.strftime('%Y%m%d%H%M')+'.npz') + if(os.path.exists(temp_path)): + temp_radar.append(temp_path) + if(len(temp_radar) == 4): + with torch.no_grad(): + R_past = torch.stack([torch.tensor(np.load(temp_radar[i])['precipitation']) for i in range(0, len(temp_radar))],dim=0) + R_past.to('cuda:1') + R_pred = fc1(R_past,num_diffusion_iters=num_diffusion_iters) + npz_data = {'precipitation':R_pred[-3]} + pred_time = name_datetime + timedelta(minutes = 180) + path_save = os.path.join(out_dir,pred_time.strftime('%Y%m%d%H%M')+'.npz') + print(path_save) + np.savez(path_save,**npz_data) +if __name__ == "__main__": + Fire(demo) +#31779 diff --git a/scripts/genpredict_2.py b/scripts/genpredict_2.py new file mode 100644 index 0000000000000000000000000000000000000000..d5d6aed354f5a64035d469536547ef00fc24dd0e --- /dev/null +++ b/scripts/genpredict_2.py @@ -0,0 +1,81 @@ +from datetime import datetime, timedelta +import glob +import os + +from fire import Fire +import h5py +from matplotlib import pyplot as plt +import numpy as np +import pathlib +import cartopy +import wradlib as wrl +import dask.array as da +import pandas as pd +import xarray as xr +import re +from datetime import datetime, timedelta +from ldcast import forecast +from ldcast.visualization import plots +import torch +MAP_PROJECTION = cartopy.crs.PlateCarree() + + +def get_date_time(name): + year=int(name[0:4]) + month=int(name[4:6]) + day=int(name[6:8]) + hour=int(name[8:10]) + minutes=int(name[10:12]) + return datetime(year,month,day,hour,minutes) + +def demo( + ldm_weights_fn="/data/data_WF/ablation/ablation_time/genforecast-radaronly-256x256-20step.pt", + autoenc_weights_fn="/home/mmhk20/weather_forecast/ldcast/models/autoenc/autoenc-32-0.01.pt", + num_diffusion_iters=50, + out_dir="/data/AI102024/test/pred_rad", + data_dir="/data/AI102024/test/rad", + t0=datetime(2020,12,31,21,40), + interval=timedelta(minutes=10), + past_timesteps=4, + crop_box=((0,640), (0,640)), + draw_border=False, + ensemble_members=1, +): + filtered_files = [] + for filename in os.listdir(data_dir): + if filename.endswith("00.npz"): + filtered_files.append(filename) + if filename.endswith("03.npz"): + filtered_files.append(filename) + sorted_files = sorted(filtered_files) + fc1 = forecast.Forecast( + ldm_weights_fn=ldm_weights_fn, + autoenc_weights_fn=autoenc_weights_fn, + future_timesteps=20, + gpu=3 + ) + index = 635 + print(len(sorted_files)) + for i in range(index, len(sorted_files)): + print(i) + print(sorted_files[i]) + name_datetime = get_date_time(sorted_files[i]) + temp_radar = [] + for i in range(0,4): + temp_time = name_datetime + timedelta(minutes = i*10 - 30) + temp_path = os.path.join(data_dir,temp_time.strftime('%Y%m%d%H%M')+'.npz') + if(os.path.exists(temp_path)): + temp_radar.append(temp_path) + if(len(temp_radar) == 4): + with torch.no_grad(): + R_past = torch.stack([torch.tensor(np.load(temp_radar[i])['precipitation']) for i in range(0, len(temp_radar))],dim=0) + R_past.to('cuda:3') + R_pred = fc1(R_past,num_diffusion_iters=num_diffusion_iters) + npz_data = {'precipitation':R_pred[-3]} + pred_time = name_datetime + timedelta(minutes = 180) + path_save = os.path.join(out_dir,pred_time.strftime('%Y%m%d%H%M')+'.npz') + print(path_save) + np.savez(path_save,**npz_data) +if __name__ == "__main__": + Fire(demo) +#31779 diff --git a/scripts/genpredict_3.py b/scripts/genpredict_3.py new file mode 100644 index 0000000000000000000000000000000000000000..9974824b31bc55a0cf38afa4e8bcd7b2b4311865 --- /dev/null +++ b/scripts/genpredict_3.py @@ -0,0 +1,81 @@ +from datetime import datetime, timedelta +import glob +import os + +from fire import Fire +import h5py +from matplotlib import pyplot as plt +import numpy as np +import pathlib +import cartopy +import wradlib as wrl +import dask.array as da +import pandas as pd +import xarray as xr +import re +from datetime import datetime, timedelta +from ldcast import forecast +from ldcast.visualization import plots +import torch +MAP_PROJECTION = cartopy.crs.PlateCarree() + + +def get_date_time(name): + year=int(name[0:4]) + month=int(name[4:6]) + day=int(name[6:8]) + hour=int(name[8:10]) + minutes=int(name[10:12]) + return datetime(year,month,day,hour,minutes) + +def demo( + ldm_weights_fn="/data/data_WF/ablation/ablation_time/genforecast-radaronly-256x256-20step.pt", + autoenc_weights_fn="/home/mmhk20/weather_forecast/ldcast/models/autoenc/autoenc-32-0.01.pt", + num_diffusion_iters=50, + out_dir="/data/AI102024/train/pred_rad", + data_dir="/data/AI102024/train/rad", + t0=datetime(2020,12,31,21,40), + interval=timedelta(minutes=10), + past_timesteps=4, + crop_box=((0,640), (0,640)), + draw_border=False, + ensemble_members=1, +): + filtered_files = [] + for filename in os.listdir(data_dir): + if filename.endswith("00.npz"): + filtered_files.append(filename) + if filename.endswith("03.npz"): + filtered_files.append(filename) + sorted_files = sorted(filtered_files) + fc1 = forecast.Forecast( + ldm_weights_fn=ldm_weights_fn, + autoenc_weights_fn=autoenc_weights_fn, + future_timesteps=20, + gpu=5 + ) + index = 19557 + print(len(sorted_files)) + for i in range(index, 20000): + print(i) + print(sorted_files[i]) + name_datetime = get_date_time(sorted_files[i]) + temp_radar = [] + for i in range(0,4): + temp_time = name_datetime + timedelta(minutes = i*10 - 30) + temp_path = os.path.join(data_dir,temp_time.strftime('%Y%m%d%H%M')+'.npz') + if(os.path.exists(temp_path)): + temp_radar.append(temp_path) + if(len(temp_radar) == 4): + with torch.no_grad(): + R_past = torch.stack([torch.tensor(np.load(temp_radar[i])['precipitation']) for i in range(0, len(temp_radar))],dim=0) + R_past.to('cuda:5') + R_pred = fc1(R_past,num_diffusion_iters=num_diffusion_iters) + npz_data = {'precipitation':R_pred[-3]} + pred_time = name_datetime + timedelta(minutes = 180) + path_save = os.path.join(out_dir,pred_time.strftime('%Y%m%d%H%M')+'.npz') + print(path_save) + np.savez(path_save,**npz_data) +if __name__ == "__main__": + Fire(demo) +#31779 diff --git a/scripts/genpredict_4.py b/scripts/genpredict_4.py new file mode 100644 index 0000000000000000000000000000000000000000..0bd0513d6ec7ab131c02bc363136a1d7880ffbd4 --- /dev/null +++ b/scripts/genpredict_4.py @@ -0,0 +1,79 @@ +from datetime import datetime, timedelta +import glob +import os + +from fire import Fire +import h5py +from matplotlib import pyplot as plt +import numpy as np +import pathlib +import cartopy +import wradlib as wrl +import dask.array as da +import pandas as pd +import xarray as xr +import re +from datetime import datetime, timedelta +from ldcast import forecast +from ldcast.visualization import plots +import torch +MAP_PROJECTION = cartopy.crs.PlateCarree() + + +def get_date_time(name): + year=int(name[0:4]) + month=int(name[4:6]) + day=int(name[6:8]) + hour=int(name[8:10]) + minutes=int(name[10:12]) + return datetime(year,month,day,hour,minutes) + +def demo( + ldm_weights_fn="/data/data_WF/ablation/ablation_time/genforecast-radaronly-256x256-20step.pt", + autoenc_weights_fn="/home/mmhk20/weather_forecast/ldcast/models/autoenc/autoenc-32-0.01.pt", + num_diffusion_iters=50, + out_dir="/data/data_WF/ablation/ablation_time/train/predict_radar", + data_dir="/data/data_WF/ablation/ablation_time/train/GT_radar", + t0=datetime(2020,12,31,21,40), + interval=timedelta(minutes=10), + past_timesteps=4, + crop_box=((0,640), (0,640)), + draw_border=False, + ensemble_members=1, +): + filtered_files = [] + for filename in os.listdir(data_dir): + if filename.endswith("00.npy") and (not filename.startswith("2020")): + filtered_files.append(filename) + if filename.endswith("03.npy") and (not filename.startswith("2020")): + filtered_files.append(filename) + sorted_files = sorted(filtered_files) + fc1 = forecast.Forecast( + ldm_weights_fn=ldm_weights_fn, + autoenc_weights_fn=autoenc_weights_fn, + future_timesteps=20, + gpu=5 + ) + index = 25002 + for i in range(index, 28500): + print(i) + print(sorted_files[i]) + name_datetime = get_date_time(sorted_files[i]) + temp_radar = [] + for i in range(0,4): + temp_time = name_datetime + timedelta(minutes = i*10 - 30) + temp_path = os.path.join(data_dir,temp_time.strftime('%Y%m%d%H%M')+'.npy') + if(os.path.exists(temp_path)): + temp_radar.append(temp_path) + if(len(temp_radar) == 4): + with torch.no_grad(): + R_past = torch.stack([torch.tensor(np.load(temp_radar[i])) for i in range(0, len(temp_radar))],dim=0) + R_past.to('cuda:5') + R_pred = fc1(R_past,num_diffusion_iters=num_diffusion_iters) + pred_time = name_datetime + timedelta(minutes = 180) + path_save = os.path.join(out_dir,pred_time.strftime('%Y%m%d%H%M')+'.npy') + np.save(path_save,R_pred[-3]) + +if __name__ == "__main__": + Fire(demo) +#31779 diff --git a/scripts/genpredict_5.py b/scripts/genpredict_5.py new file mode 100644 index 0000000000000000000000000000000000000000..09e449f78b649b1d42bc087bf37e8c35b18b460b --- /dev/null +++ b/scripts/genpredict_5.py @@ -0,0 +1,80 @@ + +from datetime import datetime, timedelta +import glob +import os + +from fire import Fire +import h5py +from matplotlib import pyplot as plt +import numpy as np +import pathlib +import cartopy +import wradlib as wrl +import dask.array as da +import pandas as pd +import xarray as xr +import re +from datetime import datetime, timedelta +from ldcast import forecast +from ldcast.visualization import plots +import torch +MAP_PROJECTION = cartopy.crs.PlateCarree() + + +def get_date_time(name): + year=int(name[0:4]) + month=int(name[4:6]) + day=int(name[6:8]) + hour=int(name[8:10]) + minutes=int(name[10:12]) + return datetime(year,month,day,hour,minutes) + +def demo( + ldm_weights_fn="/data/data_WF/ablation/ablation_time/genforecast-radaronly-256x256-20step.pt", + autoenc_weights_fn="/home/mmhk20/weather_forecast/ldcast/models/autoenc/autoenc-32-0.01.pt", + num_diffusion_iters=50, + out_dir="/data/data_WF/ablation/ablation_time/train/predict_radar", + data_dir="/data/data_WF/ablation/ablation_time/train/GT_radar", + t0=datetime(2020,12,31,21,40), + interval=timedelta(minutes=10), + past_timesteps=4, + crop_box=((0,640), (0,640)), + draw_border=False, + ensemble_members=6, +): + filtered_files = [] + for filename in os.listdir(data_dir): + if filename.endswith("00.npy") and (not filename.startswith("2020")): + filtered_files.append(filename) + if filename.endswith("03.npy") and (not filename.startswith("2020")): + filtered_files.append(filename) + sorted_files = sorted(filtered_files) + fc1 = forecast.Forecast( + ldm_weights_fn=ldm_weights_fn, + autoenc_weights_fn=autoenc_weights_fn, + future_timesteps=20, + gpu=0 + ) + index = 27600 + for i in range(index, 27750): + print(i) + print(sorted_files[i]) + name_datetime = get_date_time(sorted_files[i]) + temp_radar = [] + for i in range(0,4): + temp_time = name_datetime + timedelta(minutes = i*10 - 30) + temp_path = os.path.join(data_dir,temp_time.strftime('%Y%m%d%H%M')+'.npy') + if(os.path.exists(temp_path)): + temp_radar.append(temp_path) + if(len(temp_radar) == 4): + with torch.no_grad(): + R_past = torch.stack([torch.tensor(np.load(temp_radar[i])) for i in range(0, len(temp_radar))],dim=0) + R_past.to('cuda:0') + R_pred = fc1(R_past,num_diffusion_iters=num_diffusion_iters) + pred_time = name_datetime + timedelta(minutes = 180) + path_save = os.path.join(out_dir,pred_time.strftime('%Y%m%d%H%M')+'.npy') + np.save(path_save,R_pred[-3]) + +if __name__ == "__main__": + Fire(demo) +#18080 diff --git a/scripts/gtandpredict.py b/scripts/gtandpredict.py new file mode 100644 index 0000000000000000000000000000000000000000..b8cb3fa2bf68c6fdad5a7035461b6d2a6fe33490 --- /dev/null +++ b/scripts/gtandpredict.py @@ -0,0 +1,39 @@ +import os + +# Đường dẫn đến hai thư mục chứa file +folder1 = '/data/data_WF/ldcast_precipitation/NowcastNet/res' +folder2 = '/data/data_WF/ldcast_precipitation/test' +# Lấy danh sách các file trong mỗi thư mục +files1 = os.listdir(folder1) +files2 = os.listdir(folder2) + +# Tạo một danh sách để lưu trữ các file cùng tên +common_files = [] + +# Lặp qua từng file trong danh sách của thư mục 1 +for file1 in files1: + # Nếu file1 cũng có trong thư mục 2 + if file1 in files2: + # Thêm file1 vào danh sách các file cùng tên + common_files.append(file1) + +# Hiển thị danh sách các file cùng tên +print("Các file cùng tên trong cả hai thư mục:") +common_files.sort() +#for file_name in common_files: +# print(file_name) +j = 0 +for file1 in files1: + if file1 not in common_files: + os.remove(os.path.join(folder1, file1)) + print(file1) + j+=1 +print(j) +print("#") +i = 0 +for file2 in files2: + if file2 not in common_files: + #os.remove(os.path.join(folder2, file2)) + print(file2) + i+=1 +print(i) diff --git a/scripts/label_data.py b/scripts/label_data.py new file mode 100644 index 0000000000000000000000000000000000000000..a094c8b607ecc16255c12a1b6d74026bc5765682 --- /dev/null +++ b/scripts/label_data.py @@ -0,0 +1,33 @@ +import os +import pyart +import pathlib +from fire import Fire + + + +def label_all_items( + data_dir="/data/data_WF/NhaBe/2022", + out_dir="/data/data_WF/NhaBe/2022" +): + for subdir in os.listdir(data_dir): + subdir_path = os.path.join(data_dir, subdir) + for i in range (1,32): + if i < 10: + data_new_dir = subdir_path + '/0{}'.format(i) + out_new_dir = subdir_path + '/0{}'.format(i) + else: + data_new_dir = subdir_path + '/{}'.format(i) + out_new_dir = subdir_path + '/{}'.format(i) + pathlist = pathlib.Path(data_new_dir).rglob('*.*') + for path in pathlist: + f = os.path.join(path) + radar = pyart.io.read(f) + t = radar.time['units'].split(" ")[2][:-4] + outdir = pathlib.Path(out_new_dir) + t = outdir/t + os.rename(f, t) + # print(t) + # print(data_new_dir) + +if __name__ == "__main__": + Fire(label_all_items) diff --git a/scripts/metrics.py b/scripts/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..ba9c5ec97c5a5574895c502c181f8866127b8d1a --- /dev/null +++ b/scripts/metrics.py @@ -0,0 +1,120 @@ +import concurrent +import multiprocessing +import os + +import numpy as np + +from ldcast.analysis import crps, histogram, rank, fss +from ldcast.features import io + + +all_models=( + "mch-iters=50-res=256", + "mch-dgmr", + "mch-pysteps", + "dwd-iters=50-res=256", + "dwd-dgmr", + "dwd-pysteps", +) + + +def histogram_all(models=all_models): + os.makedirs("../results/histogram/", exist_ok=True) + for model in models: + if not model.startswith("mch"): + continue + ensembles_dir = f"../results/eval_ensembles/{model}-valid/" + hist_fn = f"../results/histogram/{model}-valid.nc" + histogram.save_histogram_for_dataset(ensembles_dir, hist_fn) + + +def prob_match_for_model(model): + parts = model.split("-") + parts = ["mch"] + parts[1:] + model = "-".join(parts) + + hist_fn = f"../results/histogram/hist-{model}-valid.nc" + (obs_hist, fc_hist, bins) = histogram.load_histogram(hist_fn) + return histogram.probability_match_timesteps(obs_hist, fc_hist, bins) + + +def crps_all(models=all_models, prob_match=True, log=False): + os.makedirs("../results/crps/", exist_ok=True) + + for model in models: + pm = prob_match_for_model(model) if prob_match else None + + ensembles_dir = f"../results/eval_ensembles/{model}/" + prefix = "logcrps" if log else "crps" + if prob_match: + prefix += "-pm" + crps_fn = f"../results/crps/{prefix}-{model}.nc" + + crps.save_crps_for_dataset( + ensembles_dir, crps_fn, log=log, preproc_fc=pm + ) + + +def ranks_all(models=all_models, prob_match=True): + os.makedirs("../results/ranks/", exist_ok=True) + + for model in models: + pm = prob_match_for_model(model) if prob_match else None + + ensembles_dir = f"../results/eval_ensembles/{model}/" + prefix = "ranks" + if prob_match: + prefix += "-pm" + ranks_fn = f"../results/ranks/{prefix}-{model}.nc" + + rank.save_ranks_for_dataset( + ensembles_dir, ranks_fn, preproc_fc=pm + ) + + +def fractions_all( + models=all_models, + thresholds=(0.1, 0.3, 1.0, 3.0, 10.0, 30.0), + prob_match=True +): + os.makedirs("../results/fractions/", exist_ok=True) + + for model in models: + pm = prob_match_for_model(model) if prob_match else None + ensembles_dir = f"../results/eval_ensembles/{model}/" + + for T in thresholds: + prefix = "fractions" + if prob_match: + prefix = prefix + "-pm" + prefix += f"-{T:.1f}" + fss_fn = f"../results/fractions/{prefix}-{model}.nc" + + fss.save_fractions_for_dataset( + ensembles_dir, fss_fn, T, preproc_fc=pm + ) + + +def rmse_ensemble_mean_batch(fn, log=False, pm=None): + print(fn) + (x, y, y_pred) = io.load_batch(fn, log=log, preproc_fc=pm) + y_pred = y_pred.mean(axis=-1) + return np.sqrt(((y-y_pred)**2).mean(axis=(1,3,4))) + + +def rmse_ensemble_mean(model="mch-iters=50-res=256", log=False, prob_match=True): + ensembles_dir = f"../results/eval_ensembles/{model}/" + files = os.listdir(ensembles_dir) + files = (os.path.join(ensembles_dir,fn) for fn in sorted(files)) + + pm = prob_match_for_model(model) if prob_match else None + + N_threads = multiprocessing.cpu_count() + tasks = [] + with concurrent.futures.ProcessPoolExecutor(N_threads) as executor: + tasks = [ + executor.submit(rmse_ensemble_mean_batch, fn, log, pm) + for fn in files + ] + rmse = [task.result() for task in tasks] + return np.concatenate(rmse, axis=0) diff --git a/scripts/train_autoenc.py b/scripts/train_autoenc.py new file mode 100644 index 0000000000000000000000000000000000000000..c2b1bafd7628ed72d2d85ca3e4e2f4109ff50c29 --- /dev/null +++ b/scripts/train_autoenc.py @@ -0,0 +1,111 @@ +import gc +import gzip +import os +import pickle + +from fire import Fire +import numpy as np +from omegaconf import OmegaConf + +from ldcast.features import batch, patches, split, transform +from ldcast.models.autoenc import encoder, training + +#file_dir = os.path.dirname(os.path.abspath(__file__)) +file_dir = os.path.dirname("/data/data_WF/ldcast_precipitation/ldcast/") + +def setup_data( + var="RZC", + batch_size=64, + sampler_file=None, + num_timesteps=8, + chunks_file="./preprocess_data/split_chunks.pkl.gz" +): + variables = { + var: { + "sources": [var], + "timesteps": np.arange(num_timesteps), + } + } + predictors = [var] # autoencoder: predictors == targets + target = var + raw_vars = [var] + raw = { + var: patches.load_all_patches( + #os.path.join(file_dir, f"../data/{var}/"), var + os.path.join(file_dir, f"./preprocess_data/{var}/"), var + ) + for var in raw_vars + } + + # Load pregenerated train/valid/test split data. + # These can be generated with features.split.get_chunks() + with gzip.open(os.path.join(file_dir, chunks_file), 'rb') as f: + chunks = pickle.load(f) + (raw, _) = split.train_valid_test_split(raw, var, chunks=chunks) + + variables[var]["transform"] = transform.default_rainrate_transform( + raw["train"][var]["scale"] + ) + + if sampler_file is None: + sampler_file = { + "train": "/data/data_WF/ldcast_precipitation/ldcast/cache/sampler_autoenc_train.pkl", + "valid": "/data/data_WF/ldcast_precipitation/ldcast/cache/sampler_autoenc_valid.pkl", + "test": "/data/data_WF/ldcast_precipitation/ldcast/cache/sampler_autoenc_test.pkl", + } + bins = np.exp(np.linspace(np.log(0.2), np.log(50), 10)) + datamodule = split.DataModule( + variables, raw, predictors, target, var, + sampling_bins=bins, batch_size=batch_size, + sampler_file=sampler_file, + valid_seed=1234, test_seed=2345 + ) + + gc.collect() + return datamodule + + +def setup_model( + model_dir=None, +): + enc = encoder.SimpleConvEncoder() + dec = encoder.SimpleConvDecoder() + (autoencoder, trainer) = training.setup_autoenc_training( + encoder=enc, + decoder=dec, + model_dir=model_dir, + ) + gc.collect() + return (autoencoder, trainer) + + +def train( + var="RZC", + batch_size=64, + sampler_file=None, + num_timesteps=8, + chunks_file="./preprocess_data/split_chunks.pkl.gz", + model_dir=None, + ckpt_path=None +): + print("Loading data...") + datamodule = setup_data( + var=var, batch_size=batch_size, sampler_file=sampler_file, + num_timesteps=num_timesteps, chunks_file=chunks_file + ) + + print("Setting up model...") + (model, trainer) = setup_model(model_dir=model_dir) + + print("Starting training...") + trainer.fit(model, datamodule=datamodule) + + +def main(config=None, **kwargs): + config = OmegaConf.load(config) if (config is not None) else {} + config.update(kwargs) + train(**config) + + +if __name__ == "__main__": + Fire(main) diff --git a/scripts/train_genforecast.py b/scripts/train_genforecast.py new file mode 100644 index 0000000000000000000000000000000000000000..8ae681ec6509d842791dd59419aa4723a66ba469 --- /dev/null +++ b/scripts/train_genforecast.py @@ -0,0 +1,129 @@ +import gc + +from fire import Fire +import torch +from omegaconf import OmegaConf + +from ldcast.models.autoenc import autoenc, encoder +from ldcast.models.genforecast import analysis, training, unet + +from train_nowcaster import setup_data + + +def setup_model( + num_timesteps=5, + model_dir="../models/test/", + autoenc_weights_fn="../models/autoenc/autoenc-32-0.01.pt", + use_obs=True, + use_nwp=False, + nwp_input_patches=4, + num_nwp_vars=9, + lr=1e-4 +): + enc = encoder.SimpleConvEncoder() + dec = encoder.SimpleConvDecoder() + autoencoder_obs = autoenc.AutoencoderKL(enc, dec) + autoencoder_obs.load_state_dict(torch.load(autoenc_weights_fn)) + + autoencoders = [] + input_patches = [] + input_size_ratios = [] + embed_dim = [] + analysis_depth = [] + if use_obs: + autoencoders.append(autoencoder_obs) + input_patches.append(1) + input_size_ratios.append(1) + embed_dim.append(128) + analysis_depth.append(4) + if use_nwp: + autoencoder_nwp = autoenc.DummyAutoencoder(width=num_nwp_vars) + autoencoders.append(autoencoder_nwp) + input_patches.append(nwp_input_patches) + input_size_ratios.append(2) + embed_dim.append(32) + analysis_depth.append(2) + + analysis_net = analysis.AFNONowcastNetCascade( + autoencoders, + input_patches=input_patches, + input_size_ratios=input_size_ratios, + train_autoenc=False, + output_patches=num_timesteps, + cascade_depth=3, + embed_dim=embed_dim, + analysis_depth=analysis_depth + ) + + model = unet.UNetModel(in_channels=autoencoder_obs.hidden_width, + model_channels=256, out_channels=autoencoder_obs.hidden_width, + num_res_blocks=2, attention_resolutions=(1,2), + dims=3, channel_mult=(1, 2, 4), num_heads=8, + num_timesteps=num_timesteps, context_ch=analysis_net.cascade_dims + ) + + (ldm, trainer) = training.setup_genforecast_training( + model, autoencoder_obs, context_encoder=analysis_net, + model_dir=model_dir, lr=lr + ) + gc.collect() + return (ldm, trainer) + + +def train( + future_timesteps=8, + use_obs=True, + use_nwp=False, + sample_shape=(4,4), + batch_size=8, + sampler=None, + ckpt_path=None, + initial_weights=None, + strict_weights=True, + model_dir=None, + lr=1e-4 +): + if sampler is None: + sampler_file = None + else: + sampler_file = { + s: f"{sampler}_{s}.pkl" for s in ["test", "train", "valid"] + } + + print("Loading data...") + datamodule = setup_data( + future_timesteps=future_timesteps, + use_obs=use_obs, + use_nwp=use_nwp, + sampler_file=sampler_file, + batch_size=batch_size, + sample_shape=sample_shape + ) + + print("Setting up model...") + (model, trainer) = setup_model( + num_timesteps=future_timesteps//4, + use_obs=use_obs, + use_nwp=use_nwp, + model_dir=model_dir, + lr=lr + ) + if initial_weights is not None: + print(f"Loading weights from {initial_weights}...") + model.load_state_dict( + torch.load(initial_weights, map_location=model.device), + strict=strict_weights + ) + + print("Starting training...") + trainer.fit(model, datamodule=datamodule, ckpt_path=ckpt_path) + + +def main(config=None, **kwargs): + config = OmegaConf.load(config) if (config is not None) else {} + config.update(kwargs) + train(**config) + + +if __name__ == "__main__": + Fire(main) diff --git a/scripts/train_nowcaster.py b/scripts/train_nowcaster.py new file mode 100644 index 0000000000000000000000000000000000000000..a27f2aecff98a90b8cd45cd1a34fef3e3de1a0fd --- /dev/null +++ b/scripts/train_nowcaster.py @@ -0,0 +1,140 @@ +from datetime import timedelta +import gc +import gzip +import os +import pickle + +import numpy as np + +from ldcast.features import batch, patches, split, transform + +#file_dir = os.path.dirname(os.path.abspath(__file__)) +file_dir = os.path.dirname("/data/data_WF/ldcast_precipitation/ldcast/") + +def setup_data( + use_obs=True, + use_nwp=False, + obs_vars=("RZC",), + nwp_vars=( + "cape", "cin", "rate-cp", "rate-tp", "t2m", + "tclw", "tcwv", "u", "v" + ), + nwp_lags=(0,12), + target_var="RZC", + batch_size=8, + past_timesteps=4, + future_timesteps=20, + timestep_secs=300, + nwp_timestep_secs=3600, + sampler_file=None, + #chunks_file="./preprocess_data/split_chunks.pkl.gz", + chunks_file="./data/split_chunks.pkl.gz", + sample_shape=(4,4) +): + target = target_var + "-T" + predictors_obs = [v + "-O" for v in obs_vars] + predictors = [] + if use_obs: + predictors += predictors_obs + if use_nwp: + predictors.append("nwp") + + variables = { + target: { + "sources": [target_var], + "timesteps": np.arange(1,future_timesteps+1), + } + } + for (var, raw_var) in zip(predictors_obs, obs_vars): + variables[var] = { + "sources": [raw_var], + "timesteps": np.arange(-past_timesteps+1,1) + } + nwp_t1 = int(np.ceil(future_timesteps*timestep_secs/nwp_timestep_secs)) + 2 + nwp_range = np.arange(nwp_t1) + variables["nwp"] = { + "sources": nwp_vars, + "timesteps": nwp_range, + "timestep_secs": nwp_timestep_secs + } + + # determine which raw variables are needed, then load them + raw_vars = set.union( + *(set(variables[v]["sources"]) for v in predictors_obs+[target]) + ) + if use_nwp: + for raw_var_base in variables["nwp"]["sources"]: + raw_vars.update(f"{raw_var_base}-{lag}" for lag in nwp_lags) + raw = { + var: patches.load_all_patches( + os.path.join(file_dir, f"./data/{var}/"), var + #os.path.join(file_dir, f"./preprocess_data/{var}/"), var + ) + for var in raw_vars + } + + # Load pregenerated train/valid/test split data. + # These can be generated with features.split.get_chunks() + with gzip.open(os.path.join(file_dir, chunks_file), 'rb') as f: + chunks = pickle.load(f) + (raw, _) = split.train_valid_test_split(raw, var, chunks=chunks) + + transform_rain = lambda: transform.default_rainrate_transform( + raw["train"]["RZC"]["scale"] + ) + transform_cape = lambda: transform.normalize_threshold( + log=True, + threshold=1.0, fill_value=1.0, + mean=1.530, std=0.859 + ) + transform_rate_tp = lambda: transform.normalize_threshold( + log=True, + threshold=1e-5, fill_value=1e-5, + mean=-3.831, std=0.650 + ) + transform_wind = lambda: transform.normalize(std=9.44) + + transforms = { + "RZC-T": transform_rain(), + "RZC-O": transform_rain(), + "cape": transform_cape(), + "cin": transform_cape(), + "rate-tp": transform_rate_tp(), + "rate-cp": transform_rate_tp(), + "t2m": transform.normalize(mean=286.069, std=7.323), + "tclw": transform.normalize_threshold( + log=True, + threshold=0.001, fill_value=0.001, + mean=-1.486, std=0.638 + ), + "tcwv": transform.normalize(std=17.307), + "u": transform_wind(), + "v": transform_wind() + } + transforms["nwp"] = transform.combine([transforms[v] for v in nwp_vars]) + for (var_name, var_data) in variables.items(): + var_data["transform"] = transforms[var_name] + + if sampler_file is None: + sampler_file = { + #"train": "../cache/sampler_nowcaster_train.pkl", + #"valid": "../cache/sampler_nowcaster_valid.pkl", + #"test": "../cache/sampler_nowcaster_test.pkl", + "train": "/data/data_WF/ldcast_precipitation/ldcast/cache/sampler_nowcaster_train.pkl", + "valid": "/data/data_WF/ldcast_precipitation/ldcast/cache/sampler_nowcaster_valid.pkl", + "test": "/data/data_WF/ldcast_precipitation/ldcast/cache/sampler_nowcaster_test.pkl", + } + bins = np.exp(np.linspace(np.log(0.2), np.log(50), 10)) + datamodule = split.DataModule( + variables, raw, predictors, target, target, + forecast_raw_vars=nwp_vars, + interval=timedelta(seconds=timestep_secs), + batch_size=batch_size, sampling_bins=bins, + time_range_sampling=(-past_timesteps+1,future_timesteps+1), + sampler_file=sampler_file, + sample_shape=sample_shape, + valid_seed=1234, test_seed=2345, + ) + + gc.collect() + return datamodule diff --git a/setup.py b/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..6b95dce1ba0d8054480bba95a900465ef762cb98 --- /dev/null +++ b/setup.py @@ -0,0 +1,39 @@ +import setuptools + +with open("README.md", "r") as fh: + long_description = fh.read() + +setuptools.setup( + name="ldcast", + version="0.0.1", + author="Jussi Leinonen", + description="Latent diffusion for generative precipitation nowcasting", + long_description=long_description, + long_description_content_type="text/markdown", + url="https://github.com/MeteoSwiss/ldcast", + packages=setuptools.find_packages(), + classifiers=[ + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", + ], + install_requires=[ + "dask", + "fire", + "einops", + "h5py", + "matplotlib", + "netCDF4", + "numba", + "numpy", + "omegaconf", + "pyshp", + "pytorch-lightning", + "scipy", + "torch", + "tqdm" + ], + extras_require={ + "benchmarks": ["tensorflow", "pysteps"] + }, + python_requires='>=3.8', +)