Upload folder using huggingface_hub

.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/

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.
+

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/<checkpoint_file>"
+```
+where `<checkpoint_file>` 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=<path_to_config_file>
+```
+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

autoenc-32-0.01.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa5ad4b8689aadbf702376e7afe5cb437ef5057675e78a8986837e8f28b3126e
+size 1617490

config/genforecast-radaronly-128x128-20step.yaml

@@ -0,0 +1 @@
+# this configuration is the default - no parameters to override!

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

environment/environment.yml

@@ -0,0 +1,4 @@
+name: ldcast
+channels:
+  - defaults
+prefix: /home/mmhk20/.conda/envs/ldcast

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

genforecast-radaronly-256x256-20step.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5fef86b78d29fde8ba66f51ae74f0d84ddc67b711fcab034a3130ea5ac7721cf
+size 5345469521

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)

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]

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

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)
+    ]

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]

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,:,:]

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,:,:]

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

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"]
+    )

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"]
+    )

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

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")

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

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
+

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))

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,
+            },
+        }

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)

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)

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

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

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

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

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

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

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
+        )
+    

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)

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

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}")

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

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

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)

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)
+

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)

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)

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)
+)

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)

models/autoenc/autoenc-32-0.01.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa5ad4b8689aadbf702376e7afe5cb437ef5057675e78a8986837e8f28b3126e
+size 1617490

models/autoenc/autoenc.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf3792a94ee4cf347ca498b10b851d58a0f2ce6b3062e7b59fec5761c7edbf24
+size 1616323

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"))