├── LICENSE ├── README.md ├── TODO.md ├── assets ├── 253171987-bd635cb3-ca3e-497c-9eb8-d032079cdb37.png └── 253768547-a488d7ed-c320-46b7-a2c5-c78bd37bcdde.png ├── data └── ICs │ ├── NX1024 │ └── 1.mat │ ├── NX128 │ ├── 1.mat │ ├── 10.mat │ ├── 11.mat │ ├── 12.mat │ ├── 13.mat │ ├── 14.mat │ ├── 15.mat │ ├── 16.mat │ ├── 17.mat │ ├── 18.mat │ ├── 19.mat │ ├── 2.mat │ ├── 20.mat │ ├── 3.mat │ ├── 4.mat │ ├── 5.mat │ ├── 6.mat │ ├── 7.mat │ ├── 8.mat │ └── 9.mat │ ├── NX2048 │ └── 1.mat │ ├── NX256 │ ├── 1.mat │ ├── 10.mat │ ├── 11.mat │ ├── 12.mat │ ├── 13.mat │ ├── 14.mat │ ├── 15.mat │ ├── 16.mat │ ├── 17.mat │ ├── 18.mat │ ├── 19.mat │ ├── 2.mat │ ├── 20.mat │ ├── 3.mat │ ├── 4.mat │ ├── 5.mat │ ├── 6.mat │ ├── 7.mat │ ├── 8.mat │ └── 9.mat │ ├── NX32 │ ├── 1.mat │ ├── 10.mat │ ├── 11.mat │ ├── 12.mat │ ├── 13.mat │ ├── 14.mat │ ├── 15.mat │ ├── 16.mat │ ├── 17.mat │ ├── 18.mat │ ├── 19.mat │ ├── 2.mat │ ├── 20.mat │ ├── 3.mat │ ├── 4.mat │ ├── 5.mat │ ├── 6.mat │ ├── 7.mat │ ├── 8.mat │ └── 9.mat │ ├── NX512 │ ├── 1.mat │ ├── 10.mat │ ├── 11.mat │ ├── 12.mat │ ├── 13.mat │ ├── 14.mat │ ├── 15.mat │ ├── 16.mat │ ├── 17.mat │ ├── 18.mat │ ├── 19.mat │ ├── 2.mat │ ├── 20.mat │ ├── 3.mat │ ├── 4.mat │ ├── 5.mat │ ├── 6.mat │ ├── 7.mat │ ├── 8.mat │ └── 9.mat │ └── NX64 │ ├── 1.mat │ ├── 10.mat │ ├── 11.mat │ ├── 12.mat │ ├── 13.mat │ ├── 14.mat │ ├── 15.mat │ ├── 16.mat │ ├── 17.mat │ ├── 18.mat │ ├── 19.mat │ ├── 2.mat │ ├── 20.mat │ ├── 3.mat │ ├── 4.mat │ ├── 5.mat │ ├── 6.mat │ ├── 7.mat │ ├── 8.mat │ └── 9.mat ├── examples ├── SGSterms.ipynb ├── SGSterms_GradientModel.ipynb ├── alignment_Tau_strainRate.ipynb ├── budget_analysis.ipynb ├── dealias.ipynb ├── derivative.ipynb ├── energy_enstrophy_transfers.ipynb ├── filter.ipynb ├── jax_vs_numpy.ipynb ├── parameter_recovery_eddy_viscosity.ipynb ├── run_solver.py ├── spectrum.ipynb ├── structure_function.ipynb └── symmetric_ifft_fft.ipynb ├── postprocess └── main.py ├── py2d ├── Py2D_solver.py ├── SGSModel.py ├── SGSterms.py ├── aposteriori_analysis.py ├── apriori_analysis.py ├── convection_conserved.py ├── convert.py ├── datamanager.py ├── dealias.py ├── derivative.py ├── eddy_viscosity_models.py ├── filter.py ├── gradient_model.py ├── initialize.py ├── readme.md ├── spectra.py ├── structure_function.py ├── util.py └── uv2tau_CNN.py └── pyproject.toml /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Karan Jakhar 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # py2d: High-Performance 2D Navier-Stokes solver in Python 2 | 3 | ## Table of contents 4 | 5 | * [Introduction](#Introduction) 6 | * [Requirements](#Requirements) 7 | * [Installation](#Installation) 8 | * [Example](#Example) 9 | * [A Note on JAX](#A-Note-on-JAX) 10 | * [Citations](#Citations) 11 | 12 | ## Introduction 13 | Py2D is a Python solver for incompressible 2-dimensional (2D) Navier-stokes equations. 14 | 15 | Py2D leverages JAX, a high-performance numerical computing library, allowing for rapid execution on GPUs while also seamlessly supporting CPUs for users who do not have access to GPUs. 16 | 17 | **Py2D features the following capabilities:** 18 | 19 | - Direct numerical simulations (DNS) for 2-dimensional (2D) turbulence, catering to a variety of systems including decaying, forced homogeneous, and beta-plane turbulence. 20 | - Large Eddy Simulations (LES) with Sub-Grid Scale (SGS) models. The compatible models include Smagorinsky (SMAG), Leith (LEITH), Dynamic Smagorinsky (DSMAG), Dynamic Leith (DLEITH), as well as gradient models - GM2, GM4, and GM6. 21 | - Coupling Neural Networks-based eddy viscosity or Sub-Grid Scale (SGS) terms with the LES solver. 22 | 23 | ## Requirements 24 | 25 | - python 3.10 26 | - [jax](https://pypi.org/project/jax/) 27 | - [jaxlib](https://pypi.org/project/jaxlib/) 28 | - [numpy](https://pypi.org/project/numpy/) 29 | 30 | ## Installation 31 | 32 | Clone the [py2d git repository](https://github.com/envfluids/py2d.git) to use the latest development version. 33 | ``` 34 | git clone https://github.com/envfluids/py2d.git 35 | ``` 36 | Then install py2d locally on your system 37 | ``` 38 | cd py2d 39 | pip install -e ./ 40 | ``` 41 | 42 | ## Example 43 | 44 | ``` 45 | from py2d.Py2D_solver import Py2D_solver 46 | 47 | # Script to call the function with the given parameters 48 | Py2D_solver(Re=20e3, # Reynolds number 49 | fkx=4, # Forcing wavenumber in x dimension 50 | fky=4, # Forcing wavenumber in y dimension 51 | alpha=0.1, # Rayleigh drag coefficient 52 | beta=0, # Coriolis parameter (Beta-plane turbulence) 53 | NX=32, # Number of grid points in x and y (Presuming a square domain) '32', '64', '128', '256', '512' 54 | forcing_filter = None, # None, "gaussian", "box" 55 | SGSModel_string='NoSGS', # SGS closure model/parametrization to use. 'NoSGS' (no closure) for DNS simulations. Available SGS models: 'SMAG', 'DSMAG', 'LEITH', 'DLEITH', 'PiOmegaGM2', 'PiOmegaGM4', 'PiOmegaGM6' 56 | eddyViscosityCoeff=0, # Coefficient for eddy viscosity models: Only used for SMAG and LEITH SGS Models 57 | dt=5e-3, # Time step 58 | dealias=True, # Dealiasing 59 | saveData=True, # Save data: The saved data directory would be printed. 60 | tSAVE=0.1, # Time interval to save data 61 | tTotal=1, # Length (total time) of simulation 62 | readTrue=False, 63 | ICnum=1, # Initial condition number: Choose between 1 to 20 64 | resumeSim=False, # start new simulation (False) or resume simulation (True) 65 | ) 66 | ``` 67 | 68 | ## A Note on JAX 69 | JAX can be installed for either CPU-only or GPU-supported environments. 70 | 71 | The default installation above will install JAX with CPU support but without GPU acceleration. 72 | 73 | #### For GPU support 74 | For instructions on installing with GPU support see [Installing JAX](https://jax.readthedocs.io/en/latest/installation.html). 75 | > Note: Note: GPU support requires a compatible NVIDIA GPU and the correct CUDA and CuDNN versions installed on your system. If you're unsure about your CUDA and CuDNN versions, consult the documentation for your GPU and the JAX installation guide for further guidance. 76 | 77 | ### If using JAX with Pytorch: 78 | 79 | > **Warning** 80 | > Combining PyTorch and JAX in the same codebase can be challenging due to their dependencies on CuDNN. It's crucial to ensure that JAX and Torch are compatible with the same version of CuDNN. JAX requires CuDNN version 8.6 or above for CUDA 11. However, it's important to verify that the version of PyTorch you are using is compiled against a compatible version of CuDNN. Mismatched versions can lead to runtime issues. 81 | 82 | > As of now, for CUDA 11, JAX works with CuDNN version 8.6 or newer. Ensure that the version of PyTorch you install is compatible with this CuDNN version. If you encounter version mismatch issues, you may need to adjust the versions of the libraries you install or consult the relevant documentation for guidance. 83 | 84 | Install a specific version of CuDNN that is compatible with both JAX and PyTorch: 85 | ``` 86 | conda install -c conda-forge cudnn=8.8.0.121 87 | ``` 88 | Install PyTorch with the appropriate CUDA toolkit version 89 | ``` 90 | conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch -c nvidia 91 | ``` 92 | Install JAX with CUDA 11 support 93 | ``` 94 | pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html 95 | ``` 96 | If your system uses environment modules, load the CUDA module (this step is system-dependent and may vary) 97 | ``` 98 | module load cuda 99 | ``` 100 | 101 | ## Citations 102 | - [Jakhar, K.](https://scholar.google.com/citations?user=buVddBgAAAAJ&hl=en), [Guan, Y.](https://gyf135.github.io/), [Mojgani, R.](https://www.rmojgani.com), [Chattopadhyay, A.](https://scholar.google.com/citations?user=wtHkCRIAAAAJ&hl=en), and [Hassanzadeh, P. 103 | ](https://scholar.google.com/citations?user=o3_eO6EAAAAJ&hl=en), 104 | [**Learning Closed-form Equations for Subgrid-scale Closures from High-fidelity Data: Promises and Challenges**](https://arxiv.org/abs/2306.05014) 105 | ```bibtex 106 | @article{jakhar2024learning, 107 | title={Learning closed-form equations for subgrid-scale closures from high-fidelity data: Promises and challenges}, 108 | author={Jakhar, Karan and Guan, Yifei and Mojgani, Rambod and Chattopadhyay, Ashesh and Hassanzadeh, Pedram}, 109 | journal={Journal of Advances in Modeling Earth Systems}, 110 | volume={16}, 111 | number={7}, 112 | pages={e2023MS003874}, 113 | year={2024}, 114 | publisher={Wiley Online Library} 115 | } 116 | ``` 117 | -------------------------------------------------------------------------------- /TODO.md: -------------------------------------------------------------------------------- 1 | # TODO 2 | 3 | - [ ] Code structured to run without requiring JAX 4 | - [ ] Code structured to run without Pytorch 5 | - [ ] Code structured for numba JIT for parallel CPU 6 | -------------------------------------------------------------------------------- /assets/253171987-bd635cb3-ca3e-497c-9eb8-d032079cdb37.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/assets/253171987-bd635cb3-ca3e-497c-9eb8-d032079cdb37.png -------------------------------------------------------------------------------- /assets/253768547-a488d7ed-c320-46b7-a2c5-c78bd37bcdde.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/assets/253768547-a488d7ed-c320-46b7-a2c5-c78bd37bcdde.png -------------------------------------------------------------------------------- /data/ICs/NX1024/1.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX1024/1.mat -------------------------------------------------------------------------------- /data/ICs/NX128/1.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/1.mat -------------------------------------------------------------------------------- /data/ICs/NX128/10.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/10.mat -------------------------------------------------------------------------------- /data/ICs/NX128/11.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/11.mat -------------------------------------------------------------------------------- /data/ICs/NX128/12.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/12.mat -------------------------------------------------------------------------------- /data/ICs/NX128/13.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/13.mat -------------------------------------------------------------------------------- /data/ICs/NX128/14.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/14.mat -------------------------------------------------------------------------------- /data/ICs/NX128/15.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/15.mat -------------------------------------------------------------------------------- /data/ICs/NX128/16.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/16.mat -------------------------------------------------------------------------------- /data/ICs/NX128/17.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/17.mat -------------------------------------------------------------------------------- /data/ICs/NX128/18.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/18.mat -------------------------------------------------------------------------------- /data/ICs/NX128/19.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/19.mat -------------------------------------------------------------------------------- /data/ICs/NX128/2.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/2.mat -------------------------------------------------------------------------------- /data/ICs/NX128/20.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/20.mat -------------------------------------------------------------------------------- /data/ICs/NX128/3.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/3.mat -------------------------------------------------------------------------------- /data/ICs/NX128/4.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/4.mat -------------------------------------------------------------------------------- /data/ICs/NX128/5.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/5.mat -------------------------------------------------------------------------------- /data/ICs/NX128/6.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/6.mat -------------------------------------------------------------------------------- /data/ICs/NX128/7.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/7.mat -------------------------------------------------------------------------------- /data/ICs/NX128/8.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/8.mat -------------------------------------------------------------------------------- /data/ICs/NX128/9.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX128/9.mat -------------------------------------------------------------------------------- /data/ICs/NX2048/1.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX2048/1.mat -------------------------------------------------------------------------------- /data/ICs/NX256/1.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/1.mat -------------------------------------------------------------------------------- /data/ICs/NX256/10.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/10.mat -------------------------------------------------------------------------------- /data/ICs/NX256/11.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/11.mat -------------------------------------------------------------------------------- /data/ICs/NX256/12.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/12.mat -------------------------------------------------------------------------------- /data/ICs/NX256/13.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/13.mat -------------------------------------------------------------------------------- /data/ICs/NX256/14.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/14.mat -------------------------------------------------------------------------------- /data/ICs/NX256/15.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/15.mat -------------------------------------------------------------------------------- /data/ICs/NX256/16.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/16.mat -------------------------------------------------------------------------------- /data/ICs/NX256/17.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/17.mat -------------------------------------------------------------------------------- /data/ICs/NX256/18.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/18.mat -------------------------------------------------------------------------------- /data/ICs/NX256/19.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/19.mat -------------------------------------------------------------------------------- /data/ICs/NX256/2.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/2.mat -------------------------------------------------------------------------------- /data/ICs/NX256/20.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/20.mat -------------------------------------------------------------------------------- /data/ICs/NX256/3.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/3.mat -------------------------------------------------------------------------------- /data/ICs/NX256/4.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/4.mat -------------------------------------------------------------------------------- /data/ICs/NX256/5.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/5.mat -------------------------------------------------------------------------------- /data/ICs/NX256/6.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/6.mat -------------------------------------------------------------------------------- /data/ICs/NX256/7.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/7.mat -------------------------------------------------------------------------------- /data/ICs/NX256/8.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/8.mat -------------------------------------------------------------------------------- /data/ICs/NX256/9.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX256/9.mat -------------------------------------------------------------------------------- /data/ICs/NX32/1.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/1.mat -------------------------------------------------------------------------------- /data/ICs/NX32/10.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/10.mat -------------------------------------------------------------------------------- /data/ICs/NX32/11.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/11.mat -------------------------------------------------------------------------------- /data/ICs/NX32/12.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/12.mat -------------------------------------------------------------------------------- /data/ICs/NX32/13.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/13.mat -------------------------------------------------------------------------------- /data/ICs/NX32/14.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/14.mat -------------------------------------------------------------------------------- /data/ICs/NX32/15.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/15.mat -------------------------------------------------------------------------------- /data/ICs/NX32/16.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/16.mat -------------------------------------------------------------------------------- /data/ICs/NX32/17.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/17.mat -------------------------------------------------------------------------------- /data/ICs/NX32/18.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/18.mat -------------------------------------------------------------------------------- /data/ICs/NX32/19.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/19.mat -------------------------------------------------------------------------------- /data/ICs/NX32/2.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/2.mat -------------------------------------------------------------------------------- /data/ICs/NX32/20.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/20.mat -------------------------------------------------------------------------------- /data/ICs/NX32/3.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/3.mat -------------------------------------------------------------------------------- /data/ICs/NX32/4.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/4.mat -------------------------------------------------------------------------------- /data/ICs/NX32/5.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/5.mat -------------------------------------------------------------------------------- /data/ICs/NX32/6.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/6.mat -------------------------------------------------------------------------------- /data/ICs/NX32/7.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/7.mat -------------------------------------------------------------------------------- /data/ICs/NX32/8.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/8.mat -------------------------------------------------------------------------------- /data/ICs/NX32/9.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX32/9.mat -------------------------------------------------------------------------------- /data/ICs/NX512/1.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/1.mat -------------------------------------------------------------------------------- /data/ICs/NX512/10.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/10.mat -------------------------------------------------------------------------------- /data/ICs/NX512/11.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/11.mat -------------------------------------------------------------------------------- /data/ICs/NX512/12.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/12.mat -------------------------------------------------------------------------------- /data/ICs/NX512/13.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/13.mat -------------------------------------------------------------------------------- /data/ICs/NX512/14.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/14.mat -------------------------------------------------------------------------------- /data/ICs/NX512/15.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/15.mat -------------------------------------------------------------------------------- /data/ICs/NX512/16.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/16.mat -------------------------------------------------------------------------------- /data/ICs/NX512/17.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/17.mat -------------------------------------------------------------------------------- /data/ICs/NX512/18.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/18.mat -------------------------------------------------------------------------------- /data/ICs/NX512/19.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/19.mat -------------------------------------------------------------------------------- /data/ICs/NX512/2.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/2.mat -------------------------------------------------------------------------------- /data/ICs/NX512/20.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/20.mat -------------------------------------------------------------------------------- /data/ICs/NX512/3.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/3.mat -------------------------------------------------------------------------------- /data/ICs/NX512/4.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/4.mat -------------------------------------------------------------------------------- /data/ICs/NX512/5.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/5.mat -------------------------------------------------------------------------------- /data/ICs/NX512/6.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/6.mat -------------------------------------------------------------------------------- /data/ICs/NX512/7.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/7.mat -------------------------------------------------------------------------------- /data/ICs/NX512/8.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/8.mat -------------------------------------------------------------------------------- /data/ICs/NX512/9.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX512/9.mat -------------------------------------------------------------------------------- /data/ICs/NX64/1.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/1.mat -------------------------------------------------------------------------------- /data/ICs/NX64/10.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/10.mat -------------------------------------------------------------------------------- /data/ICs/NX64/11.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/11.mat -------------------------------------------------------------------------------- /data/ICs/NX64/12.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/12.mat -------------------------------------------------------------------------------- /data/ICs/NX64/13.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/13.mat -------------------------------------------------------------------------------- /data/ICs/NX64/14.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/14.mat -------------------------------------------------------------------------------- /data/ICs/NX64/15.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/15.mat -------------------------------------------------------------------------------- /data/ICs/NX64/16.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/16.mat -------------------------------------------------------------------------------- /data/ICs/NX64/17.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/17.mat -------------------------------------------------------------------------------- /data/ICs/NX64/18.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/18.mat -------------------------------------------------------------------------------- /data/ICs/NX64/19.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/19.mat -------------------------------------------------------------------------------- /data/ICs/NX64/2.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/2.mat -------------------------------------------------------------------------------- /data/ICs/NX64/20.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/20.mat -------------------------------------------------------------------------------- /data/ICs/NX64/3.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/3.mat -------------------------------------------------------------------------------- /data/ICs/NX64/4.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/4.mat -------------------------------------------------------------------------------- /data/ICs/NX64/5.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/5.mat -------------------------------------------------------------------------------- /data/ICs/NX64/6.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/6.mat -------------------------------------------------------------------------------- /data/ICs/NX64/7.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/7.mat -------------------------------------------------------------------------------- /data/ICs/NX64/8.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/8.mat -------------------------------------------------------------------------------- /data/ICs/NX64/9.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/envfluids/py2d/443b30833babd852add00ddba269eeade3b6dc07/data/ICs/NX64/9.mat -------------------------------------------------------------------------------- /examples/budget_analysis.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "# Budget Analysis of 2D Turbulence" 8 | ] 9 | }, 10 | { 11 | "cell_type": "code", 12 | "execution_count": 4, 13 | "metadata": {}, 14 | "outputs": [], 15 | "source": [ 16 | "# Energy Equation\n", 17 | "\n", 18 | "import numpy as np\n", 19 | "import h5py\n", 20 | "\n", 21 | "Re = 20000\n", 22 | "fkx = 4\n", 23 | "fky = 0\n", 24 | "chi = 0.1\n", 25 | "beta = 0\n", 26 | "\n", 27 | "dealias= False\n", 28 | "\n", 29 | "\n", 30 | "dataType = 'Re20kNX1024nx4ny0r0p1' # 'Re20kNX1024nx25ny25r0p1'\n", 31 | "# filterType = 'gaussian'\n", 32 | "# Ngrid = 32\n", 33 | "# DATA_DIR = '../../data/2D_FHIT/' + dataType + '/' + filterType + '/NX' + str(Ngrid) + '/'\n", 34 | "\n", 35 | "\n", 36 | "# with h5py.File(DATA_DIR + 'PsiVor.mat', 'r') as f:\n", 37 | "# Omega_arr = f['Vor'][:].T\n", 38 | "\n", 39 | "# with h5py.File(DATA_DIR + 'S.mat', 'r') as f:\n", 40 | "# Tau11_arr = f['S1'][:].T\n", 41 | "# Tau12_arr = f['S2'][:].T\n", 42 | "# Tau22_arr = f['S3'][:].T\n", 43 | "\n", 44 | "# with h5py.File(DATA_DIR + 'Svor.mat', 'r') as f:\n", 45 | "# sigma1_arr = f['Svor1'][:].T\n", 46 | "# sigma2_arr = f['Svor2'][:].T\n", 47 | "\n", 48 | "#*******************************************************************************\n", 49 | "\n", 50 | "DATA_DIR_DNS = '../../data/2D_FHIT/' + dataType + '/DNS/train1/'\n", 51 | "with h5py.File(DATA_DIR_DNS + 'DNS1.mat', 'r') as f:\n", 52 | " Omega_arr = f['slnWorDNS'][:].T\n", 53 | "\n", 54 | "Tau11_arr = np.zeros(Omega_arr.shape)\n", 55 | "Tau12_arr = np.zeros(Omega_arr.shape)\n", 56 | "Tau22_arr = np.zeros(Omega_arr.shape)\n", 57 | "\n", 58 | "sigma1_arr = np.zeros(Omega_arr.shape)\n", 59 | "sigma2_arr = np.zeros(Omega_arr.shape)\n", 60 | "\n", 61 | "Ngrid = 1024\n", 62 | "\n", 63 | "countSnap = 0\n" 64 | ] 65 | }, 66 | { 67 | "cell_type": "code", 68 | "execution_count": 5, 69 | "metadata": {}, 70 | "outputs": [], 71 | "source": [ 72 | "from py2d.convert import Omega2Psi, Psi2UV, strain_rate\n", 73 | "from py2d.derivative import derivative\n", 74 | "from py2d.initialize import initialize_wavenumbers_rfft2, gridgen\n", 75 | "from py2d.dealias import multiply_dealias\n", 76 | "from py2d.filter import filter2D\n", 77 | "\n", 78 | "def energy_budget(Re, fkx, fky, chi, beta, Omega, Tau11, Tau12, Tau22, dealias=False):\n", 79 | " Lx, Ly = 2*np.pi, 2*np.pi\n", 80 | " Kx, Ky, Kabs, Ksq, invKsq = initialize_wavenumbers_rfft2(Ngrid, Ngrid, Lx, Ly)\n", 81 | " Lx, Ly, X, Y, dx, dy = gridgen(Lx, Ly, Ngrid, Ngrid)\n", 82 | "\n", 83 | " Psi = Omega2Psi(Omega, invKsq)\n", 84 | " U, V = Psi2UV(Psi, Kx, Ky)\n", 85 | " S11, S12, S22 = strain_rate(Psi, Kx, Ky)\n", 86 | "\n", 87 | " # Advection of Energy\n", 88 | " E = 0.5 * (multiply_dealias(U,U, dealias=dealias) + multiply_dealias(V,V, dealias=dealias))\n", 89 | " Ex = derivative(E, [1,0], Kx, Ky)\n", 90 | " Ey = derivative(E, [0,1], Kx, Ky)\n", 91 | "\n", 92 | " E_advection = U*Ex + V*Ey\n", 93 | "\n", 94 | " # Divergence of u_i S_ij\n", 95 | " u_S11 = multiply_dealias(U, S11, dealias=dealias)\n", 96 | " v_S12 = multiply_dealias(V, S12, dealias=dealias)\n", 97 | " u_S12 = multiply_dealias(U, S12, dealias=dealias)\n", 98 | " v_S22 = multiply_dealias(V, S22, dealias=dealias)\n", 99 | "\n", 100 | " u_S_divergence = derivative(u_S11 + v_S12, [1,0], Kx, Ky) + derivative(u_S12 + v_S22, [0,1], Kx, Ky)\n", 101 | "\n", 102 | " # Divergence of u_i Tau_ij\n", 103 | " u_Tau11 = multiply_dealias(U, Tau11, dealias=dealias)\n", 104 | " v_Tau12 = multiply_dealias(V, Tau12, dealias=dealias)\n", 105 | " u_Tau12 = multiply_dealias(U, Tau12, dealias=dealias)\n", 106 | " v_Tau22 = multiply_dealias(V, Tau22, dealias=dealias)\n", 107 | "\n", 108 | " u_Tau_divergence = derivative(u_Tau11 + v_Tau12, [1,0], Kx, Ky) + derivative(u_Tau12 + v_Tau22, [0,1], Kx, Ky)\n", 109 | "\n", 110 | " # Energy dissipation\n", 111 | " S11_S11 = multiply_dealias(S11,S11, dealias=dealias)\n", 112 | " S12_S12 = multiply_dealias(S12,S12, dealias=dealias)\n", 113 | " S22_S22 = multiply_dealias(S22,S22, dealias=dealias)\n", 114 | " E_dissipation = 2/Re * (S11_S11 + 2*S12_S12 + S22_S22)\n", 115 | "\n", 116 | " # Energy production\n", 117 | " Tau11_S11 = multiply_dealias(Tau11,S11, dealias=dealias)\n", 118 | " Tau12_S12 = multiply_dealias(Tau12,S12, dealias=dealias)\n", 119 | " Tau22_S22 = multiply_dealias(Tau22,S22, dealias=dealias)\n", 120 | " P_Tau = - (Tau11_S11 + 2*Tau12_S12 + Tau22_S22)\n", 121 | "\n", 122 | " # Energy Production from Forcing\n", 123 | " F1 = np.sin(fky*Y) \n", 124 | " F2 = -np.sin(fkx*X)\n", 125 | " # F1_f = filter2D(F1, filterType=filterType, coarseGrainType=None, Delta=2*Lx/Ngrid)\n", 126 | " # F2_f = filter2D(F2, filterType=filterType, coarseGrainType=None, Delta=2*Ly/Ngrid)\n", 127 | " F1_f = F1\n", 128 | " F2_f = F2\n", 129 | "\n", 130 | " u_F1_f = multiply_dealias(U, F1_f, dealias=dealias)\n", 131 | " v_F2_f = multiply_dealias(V, F2_f, dealias=dealias)\n", 132 | "\n", 133 | " F_production = u_F1_f + v_F2_f\n", 134 | "\n", 135 | " # Energy decay due to Rayleigh drag\n", 136 | " R_decay = - 2*chi * E\n", 137 | "\n", 138 | " E_total = - E_dissipation - P_Tau + F_production + R_decay\n", 139 | "\n", 140 | " return np.mean(E_advection), np.mean(u_S_divergence), np.mean(u_Tau_divergence), np.mean(E_dissipation), np.mean(P_Tau), np.mean(F_production), np.mean(R_decay), np.mean(E_total)\n", 141 | "\n", 142 | "\n", 143 | "def enstrophy_budget(Re, fkx, fky, chi, beta, Omega, Sigma1, Sigma2, dealias=False):\n", 144 | "\n", 145 | " Lx, Ly = 2*np.pi, 2*np.pi\n", 146 | " Kx, Ky, Kabs, Ksq, invKsq = initialize_wavenumbers_rfft2(Ngrid, Ngrid, Lx, Ly)\n", 147 | " Lx, Ly, X, Y, dx, dy = gridgen(Lx, Ly, Ngrid, Ngrid)\n", 148 | "\n", 149 | " Psi = Omega2Psi(Omega, invKsq)\n", 150 | " U,V = Psi2UV(Psi, Kx, Ky)\n", 151 | "\n", 152 | " # Advection of enstrophy\n", 153 | " Z = 0.5 * multiply_dealias(Omega, Omega, dealias=dealias)\n", 154 | " Zx = derivative(Z, [1,0], Kx, Ky)\n", 155 | " Zy = derivative(Z, [0,1], Kx, Ky)\n", 156 | "\n", 157 | " Z_advection = multiply_dealias(U,Zx, dealias=dealias) + multiply_dealias(V,Zy, dealias=dealias)\n", 158 | "\n", 159 | " # Divergence Omega (Grad Omega)\n", 160 | " Omega_x = derivative(Omega, [1,0], Kx, Ky)\n", 161 | " Omega_y = derivative(Omega, [0,1], Kx, Ky)\n", 162 | " Omega_Omega_x = multiply_dealias(Omega, Omega_x, dealias=dealias)\n", 163 | " Omega_Omega_y = multiply_dealias(Omega, Omega_y, dealias=dealias)\n", 164 | "\n", 165 | " Omega_divergence = 1/Re * derivative(Omega_Omega_x, [1,0], Kx, Ky) + derivative(Omega_Omega_y, [0,1], Kx, Ky)\n", 166 | "\n", 167 | " # Divergence of Omega Sigma_i\n", 168 | " Omega_Sigma1 = multiply_dealias(Omega, Sigma1, dealias=dealias)\n", 169 | " Omega_Sigma2 = multiply_dealias(Omega, Sigma2, dealias=dealias)\n", 170 | " Omega_Sigma_divergence = derivative(Omega_Sigma1, [1,0], Kx, Ky) + derivative(Omega_Sigma2, [0,1], Kx, Ky)\n", 171 | "\n", 172 | " # Enstrophy dissipation\n", 173 | " Omega_x_Omega_x = multiply_dealias(Omega_x, Omega_x, dealias=dealias)\n", 174 | " Omega_y_Omega_y = multiply_dealias(Omega_x, Omega_y, dealias=dealias)\n", 175 | "\n", 176 | " eta = 1/Re * (Omega_x_Omega_x + Omega_y_Omega_y)\n", 177 | "\n", 178 | " # Enstrophy production\n", 179 | " Pz = multiply_dealias(Sigma1, Omega_x, dealias=dealias) + multiply_dealias(Sigma2, Omega_y, dealias=dealias)\n", 180 | "\n", 181 | " # Enstrophy production from forcing\n", 182 | " f = fkx*np.cos(fkx*X) + fky*np.cos(fky*Y)\n", 183 | " F_production_ens = - multiply_dealias(f, Omega, dealias=dealias)\n", 184 | "\n", 185 | " # Enstrophy decay due to Rayleigh drag\n", 186 | " R_decay = - 2*chi * Z\n", 187 | "\n", 188 | " Z_total = - eta - Pz + F_production_ens + R_decay\n", 189 | "\n", 190 | " return np.mean(Z_advection), np.mean(Omega_divergence), np.mean(Omega_Sigma_divergence), np.mean(eta), np.mean(Pz), np.mean(F_production_ens), np.mean(R_decay), np.mean(Z_total)\n", 191 | "\n", 192 | "\n" 193 | ] 194 | }, 195 | { 196 | "cell_type": "code", 197 | "execution_count": 6, 198 | "metadata": {}, 199 | "outputs": [ 200 | { 201 | "name": "stdout", 202 | "output_type": "stream", 203 | "text": [ 204 | "****************** Energy Budget ******************\n", 205 | "E_advection = 5.5619571448506374e-18\n", 206 | "u_S_divergence = -2.0643209364124005e-17\n", 207 | "u_Tau_divergence = 0.0\n", 208 | "E_dissipation = 0.001266188265777029\n", 209 | "P_Tau = 0.0\n", 210 | "F_production = 0.20850441987235713\n", 211 | "R_decay = -2.532376561466358\n", 212 | "Total = -0.0010386264750241818\n", 213 | "****************** Enstrophy Budget ******************\n", 214 | "Z_advection = 2.2542731570318606e-16\n", 215 | "Omega_divergence = -9.544719615367515e-16\n", 216 | "Omega_Sigma_divergence = 0.0\n", 217 | "eta = 0.40985242327569765\n", 218 | "Pz = 0.0\n", 219 | "F_production = 3.336070717957713\n", 220 | "R_decay = -2.532376561466358\n", 221 | "Total = 0.3938417332156582\n" 222 | ] 223 | } 224 | ], 225 | "source": [ 226 | "E_advection_arr = []\n", 227 | "u_S_divergence_arr = []\n", 228 | "u_Tau_divergence_arr = []\n", 229 | "E_dissipation_arr = []\n", 230 | "P_Tau_arr = []\n", 231 | "F_production_arr = []\n", 232 | "R_decay_arr = []\n", 233 | "E_total_arr = []\n", 234 | "\n", 235 | "Z_advection_arr = []\n", 236 | "Omega_divergence_arr = []\n", 237 | "Omega_Sigma_divergence_arr = []\n", 238 | "eta_arr = []\n", 239 | "Pz_arr = []\n", 240 | "F_production_ens_arr = []\n", 241 | "R_decay_arr = []\n", 242 | "Z_total_arr = []\n", 243 | "\n", 244 | "for countSnap in range(20):\n", 245 | " Omega = Omega_arr[:,:,countSnap]\n", 246 | " Tau11 = Tau11_arr[:,:,countSnap]\n", 247 | " Tau12 = Tau12_arr[:,:,countSnap]\n", 248 | " Tau22 = Tau22_arr[:,:,countSnap]\n", 249 | " Sigma1 = sigma1_arr[:,:,countSnap]\n", 250 | " Sigma2 = sigma2_arr[:,:,countSnap]\n", 251 | "\n", 252 | " E_advection, u_S_divergence, u_Tau_divergence, E_dissipation, P_Tau, F_production, R_decay, E_total = energy_budget(Re, fkx, fky, chi, beta, Omega, Tau11, Tau12, Tau22, dealias=dealias)\n", 253 | "\n", 254 | " Z_advection, Omega_divergence, Omega_Sigma_divergence, eta, Pz, F_production_ens, R_decay, Z_total = enstrophy_budget(Re, fkx, fky, chi, beta, Omega, Sigma1, Sigma2, dealias=dealias)\n", 255 | "\n", 256 | " E_advection_arr.append(E_advection)\n", 257 | " u_S_divergence_arr.append(u_S_divergence)\n", 258 | " u_Tau_divergence_arr.append(u_Tau_divergence)\n", 259 | " E_dissipation_arr.append(E_dissipation)\n", 260 | " P_Tau_arr.append(P_Tau)\n", 261 | " F_production_arr.append(F_production)\n", 262 | " R_decay_arr.append(R_decay)\n", 263 | " E_total_arr.append(E_total)\n", 264 | "\n", 265 | " Z_advection_arr.append(Z_advection)\n", 266 | " Omega_divergence_arr.append(Omega_divergence)\n", 267 | " Omega_Sigma_divergence_arr.append(Omega_Sigma_divergence)\n", 268 | " eta_arr.append(eta)\n", 269 | " Pz_arr.append(Pz)\n", 270 | " F_production_ens_arr.append(F_production_ens)\n", 271 | " R_decay_arr.append(R_decay)\n", 272 | " Z_total_arr.append(Z_total)\n", 273 | "\n", 274 | "print('****************** Energy Budget ******************')\n", 275 | "print('E_advection = ', np.mean(E_advection_arr))\n", 276 | "print('u_S_divergence = ', np.mean(u_S_divergence_arr))\n", 277 | "print('u_Tau_divergence = ', np.mean(u_Tau_divergence_arr))\n", 278 | "print('E_dissipation = ', np.mean(E_dissipation_arr))\n", 279 | "print('P_Tau = ', np.mean(P_Tau_arr))\n", 280 | "print('F_production = ', np.mean(F_production_arr))\n", 281 | "print('R_decay = ', np.mean(R_decay_arr))\n", 282 | "print('Total = ', np.mean(E_total_arr))\n", 283 | "\n", 284 | "# total = -np.asarray(E_dissipation_arr) - np.asarray(P_Tau_arr) + np.asarray(F_production_arr) + np.asarray(R_decay_arr)\n", 285 | "# print('Total = ', np.mean(total))\n", 286 | "\n", 287 | "print('****************** Enstrophy Budget ******************')\n", 288 | "print('Z_advection = ', np.mean(Z_advection_arr))\n", 289 | "print('Omega_divergence = ', np.mean(Omega_divergence_arr))\n", 290 | "print('Omega_Sigma_divergence = ', np.mean(Omega_Sigma_divergence_arr))\n", 291 | "print('eta = ', np.mean(eta_arr))\n", 292 | "print('Pz = ', np.mean(Pz_arr))\n", 293 | "print('F_production = ', np.mean(F_production_ens_arr))\n", 294 | "print('R_decay = ', np.mean(R_decay_arr))\n", 295 | "print('Total = ', np.mean(Z_total_arr))\n", 296 | "\n", 297 | "# total = -np.asarray(eta_arr) - np.asarray(Pz_arr) + np.asarray(F_production_ens_arr) + np.asarray(R_decay_arr)\n", 298 | "# print('Total = ', np.mean(total))" 299 | ] 300 | } 301 | ], 302 | "metadata": { 303 | "kernelspec": { 304 | "display_name": "jax", 305 | "language": "python", 306 | "name": "python3" 307 | }, 308 | "language_info": { 309 | "codemirror_mode": { 310 | "name": "ipython", 311 | "version": 3 312 | }, 313 | "file_extension": ".py", 314 | "mimetype": "text/x-python", 315 | "name": "python", 316 | "nbconvert_exporter": "python", 317 | "pygments_lexer": "ipython3", 318 | "version": "undefined.undefined.undefined" 319 | } 320 | }, 321 | "nbformat": 4, 322 | "nbformat_minor": 2 323 | } 324 | -------------------------------------------------------------------------------- /examples/jax_vs_numpy.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "## Comparing the jax.numpy.fft with numpy.fft" 8 | ] 9 | }, 10 | { 11 | "cell_type": "code", 12 | "execution_count": 3, 13 | "metadata": {}, 14 | "outputs": [ 15 | { 16 | "name": "stdout", 17 | "output_type": "stream", 18 | "text": [ 19 | "[[1. 1.]\n", 20 | " [1. 1.]] [[1. 1.]\n", 21 | " [1. 1.]] 1.3331533165183426e-10 1.7689410108875597e-10\n", 22 | "#\n", 23 | "[[1. 1.]\n", 24 | " [1. 1.]] [[1. 1.]\n", 25 | " [1. 1.]] 3.406227136621809e-28 3.415452867676893e-28\n", 26 | "#\n", 27 | "[[1. 1.]\n", 28 | " [1. 1.]] [[1. 1.]\n", 29 | " [1. 1.]] 1.3331533163442366e-10 1.7689410107305212e-10\n" 30 | ] 31 | } 32 | ], 33 | "source": [ 34 | "import jax.numpy as jnp\n", 35 | "from jax import jit\n", 36 | "import numpy as np\n", 37 | "import numpy.testing as npt\n", 38 | "\n", 39 | "from sklearn.metrics import mean_squared_error\n", 40 | "import scipy\n", 41 | "\n", 42 | "# Generate test data\n", 43 | "np.random.seed(0) # Ensure reproducibility\n", 44 | "data = np.random.rand(256, 256) # 2D array of random numbers\n", 45 | "\n", 46 | "# Adjusted tolerance values\n", 47 | "atol = 1e-6 # Absolute tolerance\n", 48 | "rtol = 1e-6 # Relative tolerance\n", 49 | "\n", 50 | "a = np.fft.fft2(data).flatten()\n", 51 | "b = jnp.fft.fft2(data).flatten()\n", 52 | "c = scipy.fft.fft2(data).flatten()\n", 53 | "\n", 54 | "a_real = a.real.flatten()\n", 55 | "a_imag = a.imag.flatten()\n", 56 | "b_real = b.real.flatten()\n", 57 | "b_imag = b.imag.flatten()\n", 58 | "c_real = c.real.flatten()\n", 59 | "c_imag = c.imag.flatten()\n", 60 | "\n", 61 | "# print(np.fft.fft2(data) - jnp.fft.fft2(data))\n", 62 | "# print('####')\n", 63 | "# print(np.fft.fft2(data) - scipy.fft.fft2(data))\n", 64 | "\n", 65 | "print(np.corrcoef(a_real, b_real), np.corrcoef(a_imag, b_imag), mean_squared_error(a_real, b_real), mean_squared_error(a_imag, b_imag))\n", 66 | "print('#')\n", 67 | "print(np.corrcoef(a_real, c_real), np.corrcoef(a_imag, c_imag), mean_squared_error(a_real, c_real), mean_squared_error(a_imag, c_imag))\n", 68 | "print('#')\n", 69 | "print(np.corrcoef(b_real, c_real), np.corrcoef(b_imag, c_imag), mean_squared_error(b_real, c_real), mean_squared_error(b_imag, c_imag))\n", 70 | "# print(jnp.fft.fft2(data))\n", 71 | "\n" 72 | ] 73 | }, 74 | { 75 | "cell_type": "code", 76 | "execution_count": 19, 77 | "metadata": {}, 78 | "outputs": [ 79 | { 80 | "ename": "TracerArrayConversionError", 81 | "evalue": "The numpy.ndarray conversion method __array__() was called on traced array with shape float32[10000].\nThe error occurred while tracing the function numpy_jax_func at /var/folders/x6/fx3v22fd3h33fqnrs23l8_sh0000gn/T/ipykernel_78474/457597371.py:15 for jit. This concrete value was not available in Python because it depends on the value of the argument x.\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError", 82 | "output_type": "error", 83 | "traceback": [ 84 | "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", 85 | "\u001b[0;31mTracerArrayConversionError\u001b[0m Traceback (most recent call last)", 86 | "Cell \u001b[0;32mIn[19], line 24\u001b[0m\n\u001b[1;32m 22\u001b[0m \u001b[38;5;66;03m# Warm-up and compile the JAX function\u001b[39;00m\n\u001b[1;32m 23\u001b[0m _ \u001b[38;5;241m=\u001b[39m jax_func(x)\n\u001b[0;32m---> 24\u001b[0m _ \u001b[38;5;241m=\u001b[39m \u001b[43mnumpy_jax_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 26\u001b[0m \u001b[38;5;66;03m# Measure JAX function time\u001b[39;00m\n\u001b[1;32m 27\u001b[0m start_time \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n", 87 | " \u001b[0;31m[... skipping hidden 12 frame]\u001b[0m\n", 88 | "Cell \u001b[0;32mIn[19], line 17\u001b[0m, in \u001b[0;36mnumpy_jax_func\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 15\u001b[0m \u001b[38;5;129m@jit\u001b[39m\n\u001b[1;32m 16\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnumpy_jax_func\u001b[39m(x):\n\u001b[0;32m---> 17\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msin\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m \u001b[38;5;241m2\u001b[39m \u001b[38;5;241m+\u001b[39m np\u001b[38;5;241m.\u001b[39mcos(x) \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m \u001b[38;5;241m2\u001b[39m\n", 89 | "File \u001b[0;32m/opt/homebrew/anaconda3/envs/jax/lib/python3.10/site-packages/jax/_src/core.py:605\u001b[0m, in \u001b[0;36mTracer.__array__\u001b[0;34m(self, *args, **kw)\u001b[0m\n\u001b[1;32m 604\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__array__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkw):\n\u001b[0;32m--> 605\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m TracerArrayConversionError(\u001b[38;5;28mself\u001b[39m)\n", 90 | "\u001b[0;31mTracerArrayConversionError\u001b[0m: The numpy.ndarray conversion method __array__() was called on traced array with shape float32[10000].\nThe error occurred while tracing the function numpy_jax_func at /var/folders/x6/fx3v22fd3h33fqnrs23l8_sh0000gn/T/ipykernel_78474/457597371.py:15 for jit. This concrete value was not available in Python because it depends on the value of the argument x.\nSee https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError" 91 | ] 92 | } 93 | ], 94 | "source": [ 95 | "import jax.numpy as jnp\n", 96 | "from jax import jit\n", 97 | "import numpy as np\n", 98 | "import time\n", 99 | "\n", 100 | "# Define the JAX function and JIT compile it\n", 101 | "@jit\n", 102 | "def jax_func(x):\n", 103 | " return jnp.sin(x) ** 2 + jnp.cos(x) ** 2\n", 104 | "\n", 105 | "# Define the NumPy function (not JIT compiled, for comparison)\n", 106 | "def numpy_func(x):\n", 107 | " return np.sin(x) ** 2 + np.cos(x) ** 2\n", 108 | "\n", 109 | "@jit\n", 110 | "def numpy_jax_func(x):\n", 111 | " return np.sin(x) ** 2 + np.cos(x) ** 2\n", 112 | "\n", 113 | "# Generate a large NumPy array as input\n", 114 | "x = np.random.rand(10000)\n", 115 | "\n", 116 | "# Warm-up and compile the JAX function\n", 117 | "_ = jax_func(x)\n", 118 | "_ = numpy_jax_func(x)\n", 119 | "\n", 120 | "# Measure JAX function time\n", 121 | "start_time = time.time()\n", 122 | "for _ in range(500): # Run the JAX function 500 times\n", 123 | " _ = jax_func(x)\n", 124 | "jax_time = time.time() - start_time\n", 125 | "\n", 126 | "# Measure NumPy function time\n", 127 | "start_time = time.time()\n", 128 | "for _ in range(500): # Run the NumPy function 500 times\n", 129 | " _ = numpy_func(x)\n", 130 | "numpy_time = time.time() - start_time\n", 131 | "\n", 132 | "# Measure NumPy function time\n", 133 | "start_time = time.time()\n", 134 | "for _ in range(500): # Run the NumPy function 500 times\n", 135 | " _ = numpy_jax_func(x)\n", 136 | "numpy_jax_time = time.time() - start_time\n", 137 | "\n", 138 | "print(f\"JAX computation time: {jax_time} seconds\")\n", 139 | "print(f\"NumPy-JAX computation time: {numpy_jax_time} seconds\")\n", 140 | "print(f\"NumPy computation time: {numpy_time} seconds\")\n" 141 | ] 142 | } 143 | ], 144 | "metadata": { 145 | "kernelspec": { 146 | "display_name": "jax", 147 | "language": "python", 148 | "name": "python3" 149 | }, 150 | "language_info": { 151 | "codemirror_mode": { 152 | "name": "ipython", 153 | "version": 3 154 | }, 155 | "file_extension": ".py", 156 | "mimetype": "text/x-python", 157 | "name": "python", 158 | "nbconvert_exporter": "python", 159 | "pygments_lexer": "ipython3", 160 | "version": "3.10.12" 161 | } 162 | }, 163 | "nbformat": 4, 164 | "nbformat_minor": 2 165 | } 166 | -------------------------------------------------------------------------------- /examples/run_solver.py: -------------------------------------------------------------------------------- 1 | from py2d.Py2D_solver import Py2D_solver 2 | 3 | # Script to call the function with the given parameters 4 | Py2D_solver(Re=20e3, # Reynolds number 5 | fkx=4, # Forcing wavenumber in x 6 | fky=4, # Forcing wavenumber in y 7 | alpha=0.1, # Rayleigh drag coefficient 8 | beta=2.00, # Coriolis parameter 9 | NX=32, # Number of grid points in x and y '32', '64', '128', '256', '512' 10 | forcing_filter='gaussian', # Forcing filter 'None', 'gaussian', 'box' - depends on the SGSModel 11 | SGSModel_string='NoSGS', # SGS model to use 'NoSGS', 'SMAG', 'DSMAG', 'LEITH', 'DLEITH', 'PiOmegaGM2', 'PiOmegaGM4', 'PiOmegaGM6' 12 | eddyViscosityCoeff=0, # Coefficient for eddy viscosity models: SMAG and LEITH 13 | dt=5e-4, # Time step 14 | dealias=False, # Dealiasing 15 | saveData=True, # Save data 16 | tSAVE=1, # Time interval to save data 17 | tTotal=100, # Total time of simulation 18 | readTrue=False, 19 | ICnum=1, # Initial condition number: Choose between 1 to 20 20 | resumeSim=False, # start new simulation (False) or resume simulation (True) 21 | ) 22 | -------------------------------------------------------------------------------- /postprocess/main.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Created on Sat Jul 15 23:05:58 2023 5 | 6 | @author: rmojgani 7 | """ 8 | 9 | from scipy.io import loadmat, savemat 10 | 11 | 12 | from py2d.convert import * 13 | from py2d.aposteriori_analysis import * 14 | 15 | #%% 16 | NX=128 17 | Lx = 2*np.pi 18 | dx = Lx/NX 19 | #----------------- 20 | x = np.linspace(0, Lx-dx, num=NX) 21 | kx = (2*np.pi/Lx)*np.concatenate(( 22 | np.arange(0,NX/2+1,dtype=np.float64), 23 | np.arange((-NX/2+1),0,dtype=np.float64) 24 | )) 25 | # [Y,X] = np.meshgrid(x,x) 26 | # [Ky,Kx] = np.meshgrid(kx,kx) 27 | # Ksq = onp.array((Kx**2 + Ky**2)) 28 | # Ksq[0,0] = 1e16 29 | # Kabs = np.sqrt(Ksq) 30 | # invKsq = 1/Ksq 31 | 32 | Kx, Ky, Ksq = initialize_wavenumbers_2DFHIT(NX, NX, Lx, Lx) 33 | Kabs = np.sqrt(Ksq) 34 | invKsq = 1/Ksq 35 | 36 | # w1_hat = np.fft.ifft2(Omega) 37 | # signal_hat = energy_es(w1_hat, invKsq, NX, Kabs)#, invKsq, NX ) 38 | # # signal_hat = enstrophy_es(w1_hat, Kabs, NX ) 39 | # Kplot, energy, kplot_str = spectrum_angle_average_vec(signal_hat, Kabs, NX)#, kx, , invKsq) 40 | 41 | # plt.loglog(Kplot,energy) 42 | #%% 43 | import matplotlib as mpl 44 | mpl.rcParams['figure.dpi'] = 300 45 | from matplotlib import pyplot as plt 46 | 47 | 48 | #%% 49 | def spectrum_angle_average_vec(es, Kabs, NX):#, kx):#, , invKsq): 50 | ''' 51 | Angle averaged energy spectrum 52 | ''' 53 | arr_len = int(0.5*NX)# 54 | #arr_len = int(np.ceil(0.5*np.sqrt((NX*NX + NX*NX)))) 55 | kplot = np.array(range(arr_len)) 56 | eplot = np.zeros(arr_len) 57 | 58 | # spectrum for all wavenumbers 59 | for k in kplot[1:]: 60 | unmask = np.logical_and(Kabs>=(k-0.5), Kabs<(k+0.5)) 61 | #eplot[k] = k*np.sum(es[unmask]); 62 | eplot = eplot.at[k].set(k*np.sum(es[unmask])) 63 | #eplot[0]=es[0,0] 64 | # eplot = eplot.at[1].set(k*np.sum(es[unmask])) 65 | 66 | kplot_str = '\sqrt{\kappa_x^2+\kappa_y^2}' 67 | return kplot, eplot, kplot_str 68 | #%% 69 | def energy_es(w1_hat, invKsq, NX , Kabs): 70 | 71 | psi_hat = -invKsq*w1_hat 72 | es = 0.5*np.abs((np.conj(psi_hat))*w1_hat) 73 | 74 | return es 75 | #%% 76 | def enstrophy_es(w1_hat, Kabs, NX ): 77 | # March 2023 78 | es = 0.5*np.abs((np.conj(w1_hat).T)*w1_hat) 79 | 80 | return es 81 | #%% 82 | plt.figure(figsize=(12,14)) 83 | SGSModel_list = ['DLEITH_Local','DLeith']#, 'NoSGS', 'SMAG', 'DSMAG', 'LEITH', 'DLEITH', 'PiOmegaGM2', 'PiOmegaGM4', 'PiOmegaGM6']: 84 | for (SGSModel_string, icount) in zip(SGSModel_list,range(1,3,1)): 85 | 86 | # Omega = loadmat('../results/Re20k_fkx4fky0_r0.1/'+SGSModel_string+'_/NX128/dt0.0005_IC1/data/10.mat')['Omega'] 87 | Omega = loadmat('../data/ICs/NX'+str(NX)+'/1.mat')['Omega'] 88 | 89 | w1_hat = np.fft.ifft2(Omega) 90 | # signal_hat = energy_es(w1_hat, invKsq, NX, Kabs)#, invKsq, NX ) 91 | signal_hat = enstrophy_es(w1_hat, Kabs, NX ) 92 | Kplot, energy, kplot_str = spectrum_angle_average_vec(signal_hat, Kabs, NX)#, kx, , invKsq) 93 | 94 | Z_angled_average_spectra, kkx = enstrophy_angled_average_2DFHIT(Omega, spectral = False) 95 | 96 | 97 | plt.subplot(2,2,icount) 98 | plt.title(SGSModel_string) 99 | plt.contourf(Omega, cmap='bwr_r',levels=99) 100 | 101 | plt.subplot(2,2,3) 102 | plt.loglog(Kplot,energy,label=SGSModel_string) 103 | 104 | # plt.subplot(2,2,4) 105 | 106 | plt.loglog(kkx,Z_angled_average_spectra) 107 | plt.ylim([1e-6,1e0]) 108 | plt.legend() 109 | 110 | 111 | plt.show() 112 | #%% 113 | for spec_type in ['tke','enstrophy']: 114 | plt.figure(figsize=(12,14)) 115 | 116 | for (NX, icount, linewidth) in zip([128,512],range(1,3),[1,2]): 117 | 118 | print(NX) 119 | 120 | Omega = loadmat('../data/ICs/NX'+str(NX)+'/1.mat')['Omega'] 121 | 122 | Lx = 2*np.pi 123 | dx = Lx/NX 124 | #----------------- 125 | x = np.linspace(0, Lx-dx, num=NX) 126 | kx = (2*np.pi/Lx)*np.concatenate(( 127 | np.arange(0,NX/2+1,dtype=np.float64), 128 | np.arange((-NX/2+1),0,dtype=np.float64) 129 | )) 130 | # [Y,X] = np.meshgrid(x,x) 131 | # [Ky,Kx] = np.meshgrid(kx,kx) 132 | # Ksq = onp.array((Kx**2 + Ky**2)) 133 | # Ksq[0,0] = 1e16 134 | # Kabs = np.sqrt(Ksq) 135 | # invKsq = 1/Ksq 136 | 137 | Kx, Ky, Ksq = initialize_wavenumbers_2DFHIT(NX, NX, Lx, Lx) 138 | Kabs = np.sqrt(Ksq) 139 | invKsq = 1/Ksq 140 | 141 | w1_hat = np.fft.ifft2(Omega) 142 | signal_hat = energy_es(w1_hat, invKsq, NX, Kabs)#, invKsq, NX ) 143 | Kplot, energy, kplot_str = spectrum_angle_average_vec(signal_hat, Kabs, NX)#, kx, , invKsq) 144 | 145 | signal_hat = enstrophy_es(w1_hat, Kabs, NX ) 146 | Kplot, enstrophy, kplot_str = spectrum_angle_average_vec(signal_hat, Kabs, NX)#, kx, , invKsq) 147 | 148 | Omega_hat = w1_hat 149 | Omega = np.fft.ifft2(Omega_hat) 150 | Psi_hat = Omega2Psi_2DFHIT(Omega, Kx, Ky, Ksq) 151 | TKE_angled_average, kkx = TKE_angled_average_2DFHIT(Psi_hat, Omega_hat, spectral = True) 152 | Z_angled_average_spectra, kkx = enstrophy_angled_average_2DFHIT(Omega_hat, spectral = True) 153 | 154 | plt.subplot(2,2,icount) 155 | plt.title('') 156 | plt.contourf(Omega, cmap='bwr_r',levels=99) 157 | 158 | plt.subplot(2,2,3) 159 | 160 | if spec_type=='tke': 161 | plt.loglog(Kplot,energy,label='envfluids/spectra',linewidth=linewidth) 162 | elif spec_type=='enstrophy': 163 | plt.loglog(Kplot,enstrophy,label='envfluids/spectra',linewidth=linewidth) 164 | 165 | plt.legend() 166 | 167 | plt.subplot(2,2,4) 168 | if spec_type=='tke': 169 | plt.loglog(kkx,TKE_angled_average,label='py2d',linewidth=linewidth) 170 | elif spec_type=='enstrophy': 171 | plt.loglog(kkx,Z_angled_average_spectra,label='py2d') 172 | plt.xlim([0,256]) 173 | # plt.ylim([1e-9,1e1]) 174 | 175 | for pcount in [3,4]: 176 | plt.subplot(2,2,pcount) 177 | plt.ylabel(spec_type,fontsize=16) 178 | 179 | plt.legend() 180 | 181 | 182 | plt.show() -------------------------------------------------------------------------------- /py2d/Py2D_solver.py: -------------------------------------------------------------------------------- 1 | # ---------------------------------------------------------------------- 2 | # Created : Yifei Guan, Rambod Mojgani 2023 3 | # Revised : Moein Darman | Karan Jakhar May 2023 4 | # ---------------------------------------------------------------------- 5 | 6 | # 2D Turbulence Solver using Fourier-Fourier Pseudo-spectral Method 7 | # Navier-Stokes equation is in vorticity-stream function form 8 | 9 | # Import os module 10 | import os 11 | from pathlib import Path 12 | 13 | # Import Python Libraries 14 | import jax 15 | from jax import jit 16 | import numpy as nnp 17 | import jax.numpy as np 18 | from scipy.io import loadmat, savemat 19 | import time as runtime 20 | from timeit import default_timer as timer 21 | 22 | print('JAX is using ', jax.default_backend(), jax.devices()) 23 | 24 | # Import Custom Module 25 | from py2d.convection_conserved import convection_conserved, convection_conserved_dealias 26 | from py2d.convert import Omega2Psi_spectral, Psi2UV_spectral 27 | from py2d.SGSModel import * 28 | from py2d.util import regrid 29 | from py2d.filter import filter2D 30 | 31 | # from py2d.uv2tau_CNN import * 32 | 33 | from py2d.initialize import gridgen, initialize_wavenumbers_rfft2, initialize_perturbation 34 | from py2d.datamanager import gen_path, get_last_file, set_last_file, save_settings, pretty_print_table 35 | 36 | # Enable x64 Precision for Jax 37 | jax.config.update('jax_enable_x64', True) 38 | 39 | ## -------------- Initialize the kernels in JIT -------------- 40 | Omega2Psi_spectral = jit(Omega2Psi_spectral) 41 | Psi2UV_spectral = jit(Psi2UV_spectral) 42 | # prepare_data_cnn_jit = jit(prepare_data_cnn) 43 | # postproccess_data_cnn_jit = jit(postproccess_data_cnn) 44 | 45 | # Start timer 46 | startTime = timer() 47 | 48 | def Py2D_solver(Re, fkx, fky, alpha, beta, NX, forcing_filter, SGSModel_string, eddyViscosityCoeff, dt, dealias, saveData, tSAVE, tTotal, readTrue, ICnum, resumeSim): 49 | 50 | # -------------- RUN Configuration -------------- 51 | # Use random initial condition or read initialization from a file or use 52 | # readTrue = False 53 | # False: IC from file 54 | # True: Random IC (not validated) 55 | 56 | # Read the following initial condition from a file. ICnum can be from 1 to 20. 57 | # ICnumx = 5 # 1 to 20 58 | 59 | # Resume simulation from last run 60 | # resumeSim = False 61 | # True: You are ontinuing the simulation 62 | # False: You are starting the simulation from a given IC (First simulation) 63 | 64 | # Save data at the end of the simulation 65 | # saveData = True 66 | 67 | # Save snapshot at every time step 68 | # tSAVE = 0.1 69 | 70 | # Length of simulation (in time) 71 | # tTotal = 1 72 | 73 | # -------------- Geometry and mesh Parameters -------------- 74 | # Number of grid points in each direction 75 | # NX = 128 76 | 77 | # -------------- Flow specifications -------------- 78 | # Reynolds number 79 | # Re = 20e3 80 | 81 | # Time step 82 | # dt = 5e-4 83 | 84 | # Density 85 | rho = 1 86 | 87 | # Kinematic Viscosity 88 | nu = 1.0 / Re 89 | 90 | # Linear drag coefficient 91 | # alpha = 0.1 92 | 93 | # SGS Model 94 | # SGSModel_string = 'NoSGS' # SMAG DSMAG LEITH DLEITH CNN GAN 95 | 96 | # Eddy Viscosity Coefficient 97 | # eddyViscosityCoeff = 0.1 needed for SMAG and LEITH 98 | 99 | # SGS Model - PiOmega numerical scheme (Crank nicholson scheme used for time integration and eddy viscosity) 100 | PiOmega_numerical_scheme = 'E1' # E1: Euler 1st order, AB2: Adam Bashforth 2nd order 101 | 102 | # -------------- Deterministic forcing Parameters-------------- 103 | # Wavenumber in x direction 104 | # fkx = 4 105 | 106 | # Wavenumber in y direction 107 | # fky = 0 108 | 109 | # Forcing filter 110 | # forcing_filter = None # None gaussian box 111 | 112 | # -------------- Geometry and mesh Calculation -------------- 113 | 114 | # Domain length 115 | Lx = 2 * np.pi 116 | 117 | # Filter Width 118 | Delta = 2 * Lx / NX 119 | 120 | Lx, _, X, Y, dx, dx = gridgen(Lx, Lx, NX, NX) 121 | # -------------- Create the meshgrid both in physical and spectral space -------------- 122 | Kx, Ky, _, Ksq, invKsq = initialize_wavenumbers_rfft2(NX, NX, Lx, Lx) 123 | 124 | # Numpy to jax 125 | X = np.array(X) 126 | Y = np.array(Y) 127 | Kx = np.array(Kx) 128 | Ky = np.array(Ky) 129 | Ksq = np.array(Ksq) 130 | invKsq = np.array(invKsq) 131 | 132 | # -------------- Deterministic forcing Calculation -------------- 133 | 134 | # Deterministic forcing in Physical space 135 | Fk = fky * np.cos(fky * Y) + fkx * np.cos(fkx * X) 136 | 137 | # Deterministic forcing in Fourier space 138 | Fk_hat = np.fft.rfft2(Fk) 139 | 140 | # Filtering the forcing 141 | Fk_hat = filter2D(Fk_hat, filterType=forcing_filter, coarseGrainType=None, Delta=Delta, spectral=True) 142 | 143 | # -------------- RUN Configuration -------------- 144 | 145 | # Save data at every Nth iteration 146 | NSAVE = int(tSAVE / dt) 147 | 148 | # Total number of iterations 149 | maxit = int(tTotal / dt) 150 | 151 | # -------------- Directory to store data ------------------ 152 | # Snapshots of data save at the following directory 153 | SAVE_DIR, SAVE_DIR_DATA, SAVE_DIR_IC = gen_path(NX, dt, ICnum, Re, fkx, fky, alpha, beta, SGSModel_string, dealias) 154 | 155 | # Create directories if they aren't present 156 | try: 157 | os.makedirs(SAVE_DIR_DATA) 158 | os.makedirs(SAVE_DIR_IC) 159 | except OSError as error: 160 | print(error) 161 | 162 | # -------------- Print the run configuration -------------- 163 | 164 | table_flow_spec2 = [["Reynolds Number (Re)", Re], 165 | ["Deterministic Forcing Wavenumber (fkx)", fkx], 166 | ["Deterministic Forcing Wavenumber (fky)", fky], 167 | ["Forcing Filter", forcing_filter], 168 | ["Linear Drag Coefficient (alpha)", alpha], 169 | ["Beta plan coefficient (beta)", beta], 170 | ["SGS Model ", SGSModel_string], 171 | ["Eddy Viscosity Coefficient (eddyViscosityCoeff)", eddyViscosityCoeff], 172 | ["Saving Directory", SAVE_DIR_DATA]] 173 | 174 | geometry_mesh2 = [["Number of Grid Points (NX)", NX], 175 | ["Domain Length (L)", Lx], 176 | ["Mesh size (dx)", dx]] 177 | 178 | run_config2 = [["Time Step (dt)", dt], 179 | ["De-aliasing", dealias], 180 | ["Resume Simulation", resumeSim], 181 | ["Read Initialization (readTrue), If False: Will read IC from a file", readTrue], 182 | ["Initial Condition Number (ICnum)", ICnum], 183 | ["Saving Data (saveData)", saveData], 184 | ["Save data every t th timestep (tSAVE)", tSAVE], 185 | ["Save data every Nth iteration (NSAVE)", NSAVE], 186 | ["Length of simulation (tTotal)", tTotal], 187 | ["Maximum Number of Iterations (maxit)", maxit]] 188 | 189 | pretty_print_table("System Parameters", table_flow_spec2) 190 | pretty_print_table("Geometry and Mesh", geometry_mesh2) 191 | pretty_print_table("Run Configuration", run_config2) 192 | 193 | # Save the parameters to a file in writing 'w' or appending 'a' mode 194 | if resumeSim: 195 | mode = 'a' 196 | else: 197 | mode = 'w' 198 | # Open file in write mode and save parameters 199 | filename = SAVE_DIR + 'parameters.txt' 200 | with open(filename, mode) as f: 201 | 202 | empty_variable = [['', ''],['****************************************************', '****************************************************'], ['', '']] 203 | 204 | # Write each variable to a new line in the file 205 | for items in [table_flow_spec2, geometry_mesh2, run_config2, empty_variable]: 206 | for item in items: 207 | f.write(f'{item[0]}: {item[1]}\n') 208 | 209 | print("Parameters of the flow saved to saved to " + filename) 210 | 211 | # -------------- Initialization Section-------------- 212 | print("-------------- Initialization Section--------------") 213 | 214 | # -------------- Initialize PiOmega Model -------------- 215 | 216 | # PiOmega_eddyViscosity_model = SGSModel() # Initialize SGS Model 217 | PiOmega_eddyViscosity_model=SGSModel(Kx, Ky, Ksq, Delta, method=SGSModel_string, C_MODEL=eddyViscosityCoeff, dealias=dealias) 218 | # PiOmega_eddyViscosity_model.set_method(SGSModel_string) # Set SGS model to calculate PiOmega and Eddy Viscosity 219 | 220 | if SGSModel_string == 'CNN': 221 | model_path = "best_model_mcwiliams_exact.pt" 222 | cnn_model = init_model(model_type='mcwiliams', model_path=model_path) 223 | 224 | Omega0_hat, Omega1_hat, Psi0_hat, Psi1_hat, time,last_file_number_IC, last_file_number_data = initialize_conditions( 225 | NX, Kx, Ky, invKsq, readTrue, resumeSim, ICnum, SAVE_DIR_IC, SAVE_DIR_DATA ) 226 | 227 | # -------------- Main iteration loop -------------- 228 | print("-------------- Main iteration loop --------------") 229 | ## 0 meanns previous time step, 1 means current time step 230 | start_time = runtime.time() 231 | 232 | for it in range(maxit): 233 | 234 | 235 | if it == 0: 236 | U0_hat, V0_hat = Psi2UV_spectral(Psi0_hat, Kx, Ky) 237 | U1_hat, V1_hat = U0_hat, V0_hat 238 | 239 | if dealias: 240 | convec0_hat = convection_conserved_dealias(Omega0_hat, U0_hat, V0_hat, Kx, Ky) 241 | else: 242 | convec0_hat = convection_conserved(Omega0_hat, U0_hat, V0_hat, Kx, Ky) 243 | 244 | if dealias: 245 | convec1_hat = convection_conserved_dealias(Omega1_hat, U1_hat, V1_hat, Kx, Ky) 246 | else: 247 | convec1_hat = convection_conserved(Omega1_hat, U1_hat, V1_hat, Kx, Ky) 248 | 249 | # 2 Adam bash forth 250 | convec_hat = 1.5*convec1_hat - 0.5*convec0_hat 251 | 252 | diffu_hat = -Ksq*Omega1_hat 253 | 254 | PiOmega_eddyViscosity_model.update_state(Psi1_hat,Omega1_hat,U1_hat,V1_hat) 255 | PiOmega_eddyViscosity_model.calculate() 256 | 257 | 258 | PiOmega1_hat = PiOmega_eddyViscosity_model.PiOmega_hat 259 | eddyViscosity = PiOmega_eddyViscosity_model.eddy_viscosity 260 | eddyViscosityCoeff = PiOmega_eddyViscosity_model.C_MODEL 261 | 262 | # elif SGSModel_string == 'CNN': 263 | # eddyViscosity = 0.0 264 | # input_data = prepare_data_cnn_jit(Psi1_hat, Kx, Ky, Ksq) 265 | # # input_data_normalized = normalize_data(input_data) 266 | # output_normalized = PiOmega_eddyViscosity_model.calculate(cnn_model, input_data=input_data, Kx=Kx, Ky=Ky, Ksq=Ksq) 267 | # # # Diagnosis 268 | # # print("The stats of the output are: ") 269 | # # print("Mean: " + str(output_normalized.mean(axis=(1,2)))) 270 | # # print("Std: " + str(output_normalized.std(axis=(1,2)))) 271 | 272 | # # output_mean = np.array([0.0088, 5.1263e-05, 0.0108]).reshape((3,1,1)) 273 | # # output_std = np.array([0.0130, 0.0080, 0.0145]).reshape((3,1,1)) 274 | 275 | # # output_denomralized = denormalize_data(output_normalized, mean= output_mean, std= output_std) 276 | # # output_denomralized = np.zeros((3, output_normalized.shape[1], output_normalized.shape[2])) 277 | # # for i in range(3): 278 | # # updated_values = denormalize_data(output_normalized[i], mean= output_mean[i], std= output_std[i]) 279 | # # output_denomralized = output_denomralized.at[i, :, :].set(updated_values) 280 | # PiOmega1_hat = postproccess_data_cnn_jit(output_normalized[0], output_normalized[1], output_normalized[2], Kx, Ky, Ksq) 281 | # # print(np.abs(PiOmega_hat[0]).mean()) 282 | # # PiOmega_hat = PiOmegaModel.calculate() 283 | 284 | 285 | # Numerical scheme for PiOmega_hat 286 | if PiOmega_numerical_scheme == 'E1': 287 | PiOmega_hat = PiOmega1_hat 288 | 289 | elif PiOmega_numerical_scheme == 'AB2': 290 | 291 | if it ==0: 292 | PiOmega0_hat = PiOmega1_hat 293 | 294 | PiOmega_hat = 1.5*PiOmega1_hat-0.5*PiOmega0_hat 295 | 296 | # 2 Adam bash forth Crank Nicolson 297 | RHS = Omega1_hat - dt*(convec_hat) + dt*0.5*(nu+eddyViscosity)*diffu_hat - dt*(Fk_hat+PiOmega_hat) + dt*beta*V1_hat 298 | 299 | # Older version of RHS: Moein delete later 300 | # RHS = Omega1_hat + dt*(-1.5*convec1_hat + 0.5*convec0_hat) + dt*0.5*(nu+ve)*diffu_hat + dt*(Fk_hat-PiOmega_hat) 301 | 302 | Omega_hat_temp = RHS/(1+dt*alpha + 0.5*dt*(nu+eddyViscosity)*Ksq) 303 | 304 | # Replacing the previous time step with the current one 305 | Omega0_hat = Omega1_hat 306 | Omega1_hat = Omega_hat_temp 307 | convec0_hat = convec1_hat 308 | PiOmega0_hat = PiOmega1_hat 309 | 310 | # Poisson equation for Psi 311 | Psi0_hat = Psi1_hat 312 | Psi1_hat = Omega2Psi_spectral(Omega1_hat, invKsq) 313 | U1_hat, V1_hat = Psi2UV_spectral(Psi1_hat, Kx, Ky) 314 | 315 | time = time + dt 316 | 317 | if saveData and np.mod(it+1, (NSAVE)) == 0: 318 | 319 | Omega = np.real(np.fft.irfft2(Omega1_hat, s=[NX,NX])) 320 | # Psi = np.real(np.fft.irfft2(Psi1_hat, s=[NX,NX])) 321 | 322 | # Converting to numpy array 323 | Omega_cpu = nnp.array(Omega) 324 | # Psi = nnp.array(Psi) 325 | Omega0_hat_cpu = nnp.array(Omega0_hat) 326 | Omega1_hat_cpu = nnp.array(Omega1_hat) 327 | # eddyTurnoverTime = 1 / np.sqrt(np.mean(Omega ** 2)) 328 | enstrophy = 0.5 * np.mean(Omega ** 2) 329 | 330 | last_file_number_data = last_file_number_data + 1 331 | last_file_number_IC = last_file_number_IC + 1 332 | 333 | filename_data = SAVE_DIR_DATA + str(last_file_number_data) 334 | filename_IC = SAVE_DIR_IC + str(last_file_number_IC) 335 | 336 | if last_file_number_data > 2: 337 | # Remove the previous file 338 | try: 339 | os.remove(SAVE_DIR_IC + str(last_file_number_data - 2) + '.mat') 340 | except FileNotFoundError: 341 | pass 342 | 343 | try: 344 | if np.isnan(enstrophy).any(): 345 | filename = SAVE_DIR + 'unstable.txt' 346 | error_message = "eddyTurnoverTime is NaN. Stopping execution at time = " + str(time) 347 | 348 | with open(filename, 'w') as f: 349 | f.write(error_message) 350 | raise ValueError(error_message) 351 | else: 352 | savemat(filename_data + '.mat', {"Omega":Omega_cpu, "time":time}) 353 | savemat(filename_IC + '.mat', {"Omega0_hat":Omega0_hat_cpu, "Omega1_hat":Omega1_hat_cpu, "time":time, "eddyViscosity":eddyViscosity, "eddyViscosityCoeff":eddyViscosityCoeff}) 354 | 355 | except ValueError as e: 356 | print(str(e)) 357 | quit() 358 | 359 | # print('Time = {:.6f} -- Eddy Turnover Time = {:.6f} -- C = {:.4f} -- Eddy viscosity = {:.6f} ** Saved {}'.format(time, eddyTurnoverTime, eddyViscosityCoeff, eddyViscosity, filename_data)) 360 | print('Time={:.6f} -- Enstrophy={:.6f} ** file#{}'.format(time, enstrophy, last_file_number_data)) 361 | 362 | 363 | # Print elapsed time 364 | print('Total Iteration: ', it+1) 365 | endTime = timer() 366 | print('Total Time Taken: ', endTime-startTime) 367 | 368 | Omega = np.real(np.fft.irfft2(Omega1_hat, s=[NX,NX])) 369 | Omega_cpu = nnp.array(Omega) 370 | return Omega_cpu 371 | 372 | def initialize_conditions(NX, Kx, Ky, invKsq, readTrue, resumeSim, ICnum, SAVE_DIR_IC, SAVE_DIR_DATA ): 373 | 374 | if readTrue: 375 | 376 | # -------------- Initialization using pertubration -------------- 377 | w1_hat, psi_hat, psiPrevious_hat, psiCurrent_hat = initialize_perturbation(NX, Kx, Ky) 378 | time = 0.0 379 | 380 | else: 381 | 382 | if resumeSim: 383 | # Get the last file name (filenames are integers) 384 | last_file_number_data = get_last_file(SAVE_DIR_DATA) 385 | last_file_number_IC = get_last_file(SAVE_DIR_IC) 386 | 387 | # Print last file names (filenames are integers) 388 | if last_file_number_data is not None: 389 | print(f"Last data file number: {last_file_number_data}") 390 | else: 391 | print("No .mat files found") 392 | 393 | if last_file_number_IC is not None: 394 | print(f"Last IC file number: {last_file_number_IC}") 395 | else: 396 | raise ValueError("No .mat initialization files found to resume the simulation") 397 | 398 | # Load initial condition to resume simulation 399 | # Resume from the second last saved file - 400 | # the last saved file is often corrupted since the jobs stop (reach wall clocktime limit) while the file is being saved. 401 | last_file_number_data = last_file_number_data - 1 402 | last_file_number_IC = last_file_number_IC - 1 403 | 404 | data_Poi = loadmat(SAVE_DIR_IC + str(last_file_number_IC) + '.mat') 405 | Omega0_hat_cpu = data_Poi["Omega0_hat"] 406 | Omega1_hat_cpu = data_Poi["Omega1_hat"] 407 | time = data_Poi["time"] 408 | 409 | Omega_shape = Omega0_hat_cpu.shape 410 | if Omega_shape[0] == Omega_shape[1]: 411 | # If initial condition is square (fft2 fromat), convert to rfft2 format 412 | from py2d.util import fft2_to_rfft2 413 | Omega0_hat_cpu = fft2_to_rfft2(Omega0_hat_cpu) 414 | Omega1_hat_cpu = fft2_to_rfft2(Omega1_hat_cpu) 415 | else: 416 | pass 417 | 418 | # Convert numpy initialization arrays to jax array 419 | Omega0_hat = np.array(Omega0_hat_cpu) 420 | Omega1_hat = np.array(Omega1_hat_cpu) 421 | time = time[0][0] 422 | 423 | Psi0_hat = Omega2Psi_spectral(Omega0_hat, invKsq) 424 | Psi1_hat = Omega2Psi_spectral(Omega1_hat, invKsq) 425 | 426 | else: 427 | # Path of Initial Conditions 428 | 429 | # Get the absolute path to the directory 430 | base_path = Path(__file__).parent.absolute() 431 | 432 | # Construct the full path to the .mat file 433 | # Go up one directory before going into ICs 434 | if NX % 2 != 0: 435 | IC_DIR = 'data/ICs/NX' + str(NX-1) + '/' 436 | else: 437 | IC_DIR = 'data/ICs/NX' + str(NX) + '/' 438 | 439 | IC_filename = str(ICnum) + '.mat' 440 | file_path = os.path.join(base_path, "..", IC_DIR, IC_filename) 441 | 442 | # Resolve the '..' to compute the actual directory 443 | file_path = Path(file_path).resolve() 444 | 445 | # -------------- Loading Initial condition (***) -------------- 446 | 447 | data_Poi = loadmat(file_path) 448 | Omega1 = data_Poi["Omega"] 449 | if NX % 2 != 0: 450 | Omega1 = regrid(Omega1, NX, NX) 451 | 452 | Omega1_hat = np.fft.rfft2(Omega1) 453 | Omega0_hat = Omega1_hat 454 | Psi1_hat = Omega2Psi_spectral(Omega1_hat, invKsq) 455 | Psi0_hat = Psi1_hat 456 | time = 0.0 457 | 458 | # Get the last file name (filenames are integers) 459 | last_file_number_data = get_last_file(SAVE_DIR_DATA) 460 | last_file_number_IC = get_last_file(SAVE_DIR_IC) 461 | 462 | # Set last File numbers 463 | if last_file_number_data is None: 464 | print(f"Last data file number: {last_file_number_data}") 465 | last_file_number_data = 0 466 | print("Updated last data file number to " + str(last_file_number_data)) 467 | else: 468 | raise ValueError("Data already exists in the results folder for this case, either resume the simulation (resumeSim = True) or delete data to start a new simulation") 469 | 470 | if last_file_number_IC is None: 471 | print(f"Last IC file number: {last_file_number_IC}") 472 | last_file_number_IC = 0 473 | print("Updated last IC file number to " + str(last_file_number_IC)) 474 | else: 475 | raise ValueError("Data already exists in the results folder for this case, either resume the simulation (resumeSim = True) or delete data to start a new simulation") 476 | 477 | return Omega0_hat, Omega1_hat, Psi0_hat, Psi1_hat, time,last_file_number_IC, last_file_number_data 478 | 479 | if __name__ == '__main__': 480 | import sys 481 | sys.path.append('examples') 482 | sys.path.append('py2d') 483 | sys.path.append('.') 484 | #SGSModel_list = ['NoSGS', 'PiOmegaGM2', 'PiOmegaGM4', 'PiOmegaGM6'] 485 | # SGSModel_list = ['SMAG','DSMAG','DSMAG_tau_Local','DSMAG_sigma_Local'] 486 | SGSModel_list = ['DSMAG_tau_Local_LocalS','DSMAG_sigma_Local_LocalS'] 487 | #SGSModel_list = [ 'LEITH', 'DLEITH', DLEITH_tau_Local', 'DLEITH_sigma_Local'] 488 | for SGSModel_string in SGSModel_list: 489 | # Script to call the function with the given parameters 490 | Py2D_solver(Re=20e3, # Reynolds number 491 | fkx=4, # Forcing wavenumber in x 492 | fky=0, # Forcing wavenumber in y 493 | alpha=0.1, # Rayleigh drag coefficient 494 | beta=20, # Coriolis parameter 495 | NX=128, # Number of grid points in x and y '32', '64', '128', '256', '512' 496 | SGSModel_string=SGSModel_string, # SGS model to use 'NoSGS', 'SMAG', 'DSMAG', 'LEITH', 'DLEITH', 'PiOmegaGM2', 'PiOmegaGM4', 'PiOmegaGM6' 497 | eddyViscosityCoeff=0.17, # Coefficient for eddy viscosity models: SMAG and LEITH 498 | dt=5e-4, # Time step 499 | saveData=True, # Save data 500 | dealias=True, # dealias 501 | tSAVE=1.0, # Time interval to save data 502 | tTotal=10.0, # Total time of simulation 503 | readTrue=False, 504 | ICnum=1, # Initial condition number: Choose between 1 to 20 505 | resumeSim=False, # tart new simulation (False) or resume simulation (True) 506 | ) 507 | -------------------------------------------------------------------------------- /py2d/aposteriori_analysis.py: -------------------------------------------------------------------------------- 1 | # ---------------------------------------------------------------------- 2 | # Created : Karan Jakhar May 2023 3 | # ---------------------------------------------------------------------- 4 | 5 | import numpy as np 6 | 7 | from py2d.convert import strain_rate 8 | from py2d.derivative import derivative 9 | 10 | def eddyTurnoverTime(Omega): 11 | """ 12 | Compute eddy turnover time for 2D_FHIT using Omega. 13 | 14 | Args: 15 | A (ndarray): 2D array of Omega U. 16 | definition (str): Optional string to define eddy turnover time. Default is 'Enstrophy'. 17 | Possible values: 'Enstrophy', 'Omega', 'Velocity' 18 | 19 | Returns: 20 | float: Eddy turnover time. 21 | """ 22 | eddyTurnoverTime = 1 / np.sqrt(np.mean(Omega ** 2)) 23 | return eddyTurnoverTime 24 | 25 | def energyDissipationRate(Psi, Re, Kx, Ky, spectral=False): 26 | """ 27 | Compute energy dissipation rate for 2D_FHIT using Omega. 28 | 29 | Args: 30 | A (ndarray): 2D array of Omega U. 31 | Re (float): Reynolds number. 32 | 33 | Returns: 34 | float: Energy dissipation rate. 35 | 36 | Note: 37 | The energy dissipation rate is computed in Eq.16 of [1]. 38 | [1] Buaria, D., & Sreenivasan, K. R. (2023). 39 | Forecasting small-scale dynamics of fluid turbulence using deep neural networks. 40 | Proceedings of the National Academy of Sciences, 120(30), e2305765120. 41 | https://www.pnas.org/doi/abs/10.1073/pnas.2305765120 42 | """ 43 | N = Psi.shape[0] 44 | 45 | # Compute strain rate tensor 46 | if spectral: 47 | S11_hat, S12_hat, S22_hat = strain_rate(Psi, Kx, Ky, spectral=True) 48 | S11, S12, S22 = np.fft.irfft2(S11_hat, s=[N,N]), np.fft.irfft2(S12_hat, s=[N,N]), np.fft.irfft2(S22_hat, s=[N,N]) 49 | else: 50 | S11, S12, S22 = strain_rate(Psi, Kx, Ky, spectral=False) 51 | 52 | # Compute energy dissipation rate 53 | energyDissipationRate = (2/Re) * np.mean(S11 ** 2 + 2 * S12 ** 2 + S22 ** 2) 54 | 55 | return energyDissipationRate 56 | 57 | 58 | def enstrophyDissipationRate(Omega, Re, Kx, Ky, spectral=False): 59 | """ 60 | Compute enstrophy dissipation rate for 2D_FHIT using Omega. 61 | 62 | Args: 63 | A (ndarray): 2D array of Omega U. 64 | Re (float): Reynolds number. 65 | 66 | Returns: 67 | float: Enstrophy dissipation rate. 68 | 69 | Note: 70 | The energy dissipation rate is computed in Eq.11 of [1]. 71 | [1] Alexakis, A., & Doering, C. R. (2006). 72 | Energy and enstrophy dissipation in steady state 2d turbulence. 73 | Physics letters A, 359(6), 652-657. 74 | https://doi.org/10.1016/j.physleta.2006.07.048 75 | """ 76 | N = Omega.shape[0] 77 | 78 | # Compute vorticity gradient 79 | if spectral: 80 | Omega_hat = np.fft.irfft2(Omega, s=[N,N]) 81 | Omegax_hat = derivative(Omega_hat, [1,0], Kx=Kx, Ky=Ky, spectral=True) 82 | Omegay_hat = derivative(Omega_hat, [0,1], Kx=Kx, Ky=Ky, spectral=True) 83 | Omegax, Omegay = np.fft.irfft2(Omegax_hat, s=[N,N]), np.fft.irfft2(Omegay_hat, s=[N,N]) 84 | 85 | else: 86 | Omegax = derivative(Omega, [1,0], Kx=Kx, Ky=Ky, spectral=False) 87 | Omegay = derivative(Omega, [0,1], Kx=Kx, Ky=Ky, spectral=False) 88 | 89 | # Compute enstrophy dissipation rate 90 | enstrophyDissipationRate = (1/Re) * np.mean(Omegax ** 2 + Omegay ** 2) 91 | 92 | return enstrophyDissipationRate -------------------------------------------------------------------------------- /py2d/apriori_analysis.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from py2d.derivative import derivative 3 | from py2d.convert import Tau2AnisotropicTau, strain_rate 4 | from py2d.initialize import initialize_wavenumbers_rfft2 5 | from py2d.dealias import multiply_dealias 6 | from py2d.util import eig_vec_2D 7 | 8 | def energy(Psi, Omega, spectral = False, dealias = False): 9 | '''Calculates energy as the mean of 0.5 * Psi * Omega. 10 | 11 | Args: 12 | Psi (np.ndarray): The stream function matrix. 13 | Omega (np.ndarray): The vorticity matrix. 14 | spectral (bool): Whether to perform inverse Fast Fourier Transform on Psi and Omega. Default is False. 15 | 16 | Returns: 17 | float: The calculated energy. 18 | ''' 19 | N = Psi.shape[0] 20 | 21 | if spectral: 22 | Psi = np.fft.irfft2(Psi, s=[N, N]) 23 | Omega = np.fft.irfft2(Omega, s=[N, N]) 24 | 25 | PsiOmega = multiply_dealias(Psi, Omega, dealias=dealias) 26 | 27 | energy = np.mean(0.5 * PsiOmega) 28 | return energy 29 | 30 | 31 | def enstrophy(Omega, spectral = False, dealias=False): 32 | '''Calculates enstrophy as the mean of 0.5 * Omega * Omega. 33 | 34 | Args: 35 | Omega (np.ndarray): The vorticity matrix. 36 | spectral (bool): Whether to perform inverse Fast Fourier Transform on Omega. Default is False. 37 | 38 | Returns: 39 | float: The calculated enstrophy. 40 | ''' 41 | N = Omega.shape[0] 42 | 43 | if spectral: 44 | Omega = np.fft.irfft2(Omega, s=[N, N]) 45 | 46 | OmegaOmega = multiply_dealias(Omega, Omega, dealias=dealias) 47 | 48 | enstrophy = np.mean(0.5 * OmegaOmega) 49 | return enstrophy 50 | 51 | 52 | def energyTransfer(U, V, Tau11, Tau12, Tau22, Kx, Ky, dealias = False): 53 | """ 54 | Energy transfer of 2D_FHIT using SGS stress 55 | Input is single snapshot (N x N matrix) 56 | 57 | Inputs: 58 | U,V: Velocities 59 | Tau11, Tau12, Tau22: SGS stress 60 | 61 | Output: 62 | PTau: energy transfer 63 | 64 | Note: Dealias[a,b,c] != Dealias[Dealias[a,b],c] 65 | Dealiasing energy transfer directly may not render correct result and you may need to explore how Tau is caluclated. 66 | Considering Tau is usually composed fo multiplying 2 fields. 67 | """ 68 | 69 | Ux = derivative(U, [1,0], Kx=Kx, Ky=Ky, spectral=False) 70 | Uy = derivative(U, [0,1], Kx=Kx, Ky=Ky, spectral=False) 71 | Vx = derivative(V, [1,0], Kx=Kx, Ky=Ky, spectral=False) 72 | 73 | Tau11Ux = multiply_dealias(Tau11, Ux, dealias=dealias) 74 | Tau22Ux = multiply_dealias(Tau22, Ux, dealias=dealias) 75 | Tau12Uy = multiply_dealias(Tau12, Uy, dealias=dealias) 76 | Tau12Vx = multiply_dealias(Tau12, Vx, dealias=dealias) 77 | 78 | PTau = -Tau11Ux + Tau22Ux - Tau12Uy - Tau12Vx 79 | 80 | return PTau 81 | 82 | def energyTransfer_domain_average(Psi, PiOmega, dealias=False): 83 | """ 84 | Domain Average of Energy transfer of 2D_FHIT using SGS stress 85 | Input is single snapshot (N x N matrix) 86 | 87 | Inputs: 88 | Psi: Stream function 89 | PiOmega: SGS term 90 | 91 | Output: 92 | PTau: energy transfer - domain average 93 | 94 | Note: Dealias[a,b,c] != Dealias[Dealias[a,b],c] 95 | Dealiasing energy transfer directly may not render correct result and you may need to explore how Tau is caluclated. 96 | Considering Tau is usually composed fo multiplying 2 fields. 97 | """ 98 | 99 | PTau = np.mean(multiply_dealias(Psi, PiOmega, dealias=dealias)) 100 | 101 | return PTau 102 | 103 | def enstrophyTransfer(Omega, Sigma1, Sigma2, Kx, Ky, dealias = False): 104 | """ 105 | Enstrophy transfer of 2D_FHIT using SGS vorticity stress 106 | 107 | Inputs: 108 | Omega: Vorticity 109 | Sigma1, Sigma2: SGS vorticity stress 110 | 111 | Output: 112 | PZ: enstrophy transfer 113 | 114 | Note: Dealias[a,b,c] != Dealias[Dealias[a,b],c] 115 | Dealiasing enstrophy transfer directly may not render correct result and you may need to explore how Sigma is caluclated. 116 | Considering Sigma is usually composed fo multiplying 2 fields. 117 | """ 118 | 119 | Omegax = derivative(Omega, [1,0], Kx=Kx, Ky=Ky, spectral=False) 120 | Omegay = derivative(Omega, [0,1], Kx=Kx, Ky=Ky, spectral=False) 121 | 122 | Sigma1Omegax = multiply_dealias(Sigma1, Omegax, dealias=dealias) 123 | Sigma2Omegay = multiply_dealias(Sigma2, Omegay, dealias=dealias) 124 | 125 | PZ = -Sigma1Omegax - Sigma2Omegay 126 | 127 | return PZ 128 | 129 | def enstrophyTransfer_domain_average(Omega, PiOmega, dealias=False): 130 | """ 131 | Domain Average of Enstrophy transfer of 2D_FHIT using SGS vorticity stress 132 | 133 | Inputs: 134 | Omega: Vorticity 135 | PiOmega: SGS term 136 | 137 | Output: 138 | PZ: enstrophy transfer - domain average 139 | 140 | Note: Dealias[a,b,c] != Dealias[Dealias[a,b],c] 141 | Dealiasing enstrophy transfer directly may not render correct result and you may need to explore how Sigma is caluclated. 142 | Considering Sigma is usually composed fo multiplying 2 fields. 143 | """ 144 | 145 | PZ = np.mean(multiply_dealias(Omega, PiOmega, dealias=dealias)) 146 | 147 | return PZ 148 | 149 | def angle_Tau_strainRate(Tau11, Tau12, Tau22, Psi, anisotropic=False, dealias=False): 150 | """ 151 | Calculate the angle between the stress tensor and the strain rate tensor. 152 | 153 | This function takes the components of the stress tensor and the strain rate tensor and computes the angle between them. 154 | 155 | Parameters: 156 | Tau11 (np.ndarray): A 2D array representing the Tau11 component of the stress tensor. 157 | Tau12 (np.ndarray): A 2D array representing the Tau12 component of the stress tensor. 158 | Tau22 (np.ndarray): A 2D array representing the Tau22 component of the stress tensor. 159 | Psi (np.ndarray): A 2D array representing the Psi component of the strain rate tensor. 160 | anisotropic (bool): A boolean flag to indicate whether the material is anisotropic. 161 | 162 | Returns: 163 | np.ndarray: A 2D array containing the angle between the stress tensor and the strain rate tensor at each point. 164 | """ 165 | 166 | if anisotropic: 167 | Tauu11r = Tau11 168 | Tau12r = Tau12 169 | Tau22r = Tau22 170 | else: 171 | Tauu11r, Tau12r, Tau22r = Tau2AnisotropicTau(Tau11, Tau12, Tau22) 172 | 173 | # Calculate the components of the strain rate tensor 174 | nx, ny = Tau11.shape 175 | Lx, Ly = 2*np.pi, 2*np.pi 176 | Kx, Ky, _, _, _ = initialize_wavenumbers_rfft2(nx, ny, Lx, Ly) 177 | 178 | S11, S12, S22 = strain_rate(Psi, Kx, Ky) 179 | 180 | # Calculate the angle between the stress tensor and the strain rate tensor 181 | Tau_S = multiply_dealias(Tauu11r, S11, dealias=dealias) + 2*multiply_dealias( 182 | Tau12r, S12, dealias=dealias) + multiply_dealias(Tau22r, S22, dealias=dealias) 183 | Tau_Tau = multiply_dealias(Tauu11r, Tauu11r, dealias=dealias) + 2*multiply_dealias( 184 | Tau12r, Tau12r, dealias=dealias) + multiply_dealias(Tau22r, Tau22r, dealias=dealias) 185 | S_S = multiply_dealias(S11, S11, dealias=dealias) + 2*multiply_dealias( 186 | S12, S12, dealias=dealias) + multiply_dealias(S22, S22, dealias=dealias) 187 | 188 | angle = np.arccos(Tau_S/(np.sqrt(Tau_Tau)*np.sqrt(S_S))) 189 | 190 | # Calculate the angle between eigenvectors 191 | S_eigVec1, S_eigVec2, _, _ = eig_vec_2D(S11, S12, S12, S22) 192 | Tau_eigVec1, Tau_eigVec2, _, _ = eig_vec_2D(Tauu11r, Tau12r, Tau12r, Tau22r) 193 | 194 | # Dot product between eigenvectors 195 | Tau_S_eigVec1 = np.sum(Tau_eigVec1*S_eigVec1, axis=1) 196 | Tau_S_eigVec2 = np.sum(Tau_eigVec2*S_eigVec2, axis=1) 197 | 198 | Tau_Tau_eigVec1 = np.sum(Tau_eigVec1*Tau_eigVec1, axis=1) 199 | Tau_Tau_eigVec2 = np.sum(Tau_eigVec2*Tau_eigVec2, axis=1) 200 | 201 | S_S_eigVec1 = np.sum(S_eigVec1*S_eigVec1, axis=1) 202 | S_S_eigVec2 = np.sum(S_eigVec2*S_eigVec2, axis=1) 203 | 204 | angle_eigVec1 = np.arccos(Tau_S_eigVec1/(np.sqrt(Tau_Tau_eigVec1)*np.sqrt(S_S_eigVec1))) 205 | angle_eigVec2 = np.arccos(Tau_S_eigVec2/(np.sqrt(Tau_Tau_eigVec2)*np.sqrt(S_S_eigVec2))) 206 | 207 | return angle, angle_eigVec1, angle_eigVec2 208 | 209 | -------------------------------------------------------------------------------- /py2d/convection_conserved.py: -------------------------------------------------------------------------------- 1 | import jax.numpy as jnp 2 | from jax import jit 3 | 4 | from py2d.dealias import multiply_dealias_spectral_jit 5 | 6 | @jit 7 | def convection_conserved_dealias(Omega1_hat, U1_hat, V1_hat, Kx, Ky): 8 | 9 | # Convservative form 10 | # U1_hat = (1.j) * Ky * Psi1_hat 11 | # V1_hat = -(1.j) * Kx * Psi1_hat 12 | # U1 = jnp.real(jnp.fft.irfft2(U1_hat, s=[N,N])) 13 | # V1 = jnp.real(jnp.fft.irfft2(V1_hat, s=[N,N])) 14 | # Omega1 = jnp.real(jnp.fft.irfft2(Omega1_hat, s=[N,N])) 15 | 16 | # dealiasing 17 | U1Omega1_hat = multiply_dealias_spectral_jit(U1_hat, Omega1_hat) 18 | V1Omega1_hat = multiply_dealias_spectral_jit(V1_hat, Omega1_hat) 19 | 20 | conu1 = (1.j) * Kx * U1Omega1_hat 21 | conv1 = (1.j) * Ky * V1Omega1_hat 22 | convec_hat = conu1 + conv1 23 | 24 | # Non-conservative form 25 | Omega1x_hat = (1.j) * Kx * Omega1_hat 26 | Omega1y_hat = (1.j) * Ky * Omega1_hat 27 | 28 | # dealiasing 29 | U1Omega1x_hat = multiply_dealias_spectral_jit(U1_hat, Omega1x_hat) 30 | V1Omega1y_hat = multiply_dealias_spectral_jit(V1_hat, Omega1y_hat) 31 | 32 | conu1 = U1Omega1x_hat 33 | conv1 = V1Omega1y_hat 34 | convecN_hat = conu1 + conv1 35 | 36 | convec_hat = 0.5 * (convec_hat + convecN_hat) 37 | 38 | return convec_hat 39 | 40 | @jit 41 | def convection_conserved(Omega1_hat, U1_hat, V1_hat, Kx, Ky): 42 | 43 | # Convservative form 44 | # U1_hat = (1.j) * Ky * Psi1_hat 45 | # V1_hat = -(1.j) * Kx * Psi1_hat 46 | 47 | N = Omega1_hat.shape[0] 48 | 49 | U1 = jnp.fft.irfft2(U1_hat, s=[N,N]) 50 | V1 = jnp.fft.irfft2(V1_hat, s=[N,N]) 51 | Omega1 = jnp.fft.irfft2(Omega1_hat, s=[N,N]) 52 | 53 | U1Omega1 = U1*Omega1 54 | V1Omega1 = V1*Omega1 55 | 56 | conu1 = (1.j) * Kx * jnp.fft.rfft2(U1Omega1) 57 | conv1 = (1.j) * Ky * jnp.fft.rfft2(V1Omega1) 58 | convec_hat = conu1 + conv1 59 | 60 | # Non-conservative form 61 | Omega1x_hat = (1.j) * Kx * Omega1_hat 62 | Omega1y_hat = (1.j) * Ky * Omega1_hat 63 | 64 | Omega1x = jnp.fft.irfft2(Omega1x_hat, s=[N,N]) 65 | Omega1y = jnp.fft.irfft2(Omega1y_hat, s=[N,N]) 66 | 67 | U1Omega1x = U1*Omega1x 68 | V1Omega1y = V1*Omega1y 69 | 70 | conu1 = jnp.fft.rfft2(U1Omega1x) 71 | conv1 = jnp.fft.rfft2(V1Omega1y) 72 | convecN_hat = conu1 + conv1 73 | 74 | convec_hat = 0.5 * (convec_hat + convecN_hat) 75 | 76 | return convec_hat -------------------------------------------------------------------------------- /py2d/datamanager.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Created on Fri Jul 7 18:10:24 2023 5 | 6 | @author: rm99 7 | """ 8 | import os 9 | import glob 10 | # -------------------------------------------------------------------------- 11 | #def load_IC(): 12 | 13 | # -------------------------------------------------------------------------- 14 | def gen_path(NX, dt, ICnum, Re, 15 | fkx, fky, alpha, beta, SGSModel_string, dealias): 16 | # Snapshots of data save at the following directory 17 | if dealias: 18 | dataType_DIR = 'Re' + str(int(Re)) + '_fkx' + str(fkx) + 'fky' + str(fky) + '_r' + str(alpha) + '_b' + str(beta) + '/' 19 | else: 20 | dataType_DIR = 'Re' + str(int(Re)) + '_fkx' + str(fkx) + 'fky' + str(fky) + '_r' + str(alpha) + '_b' + str(beta) + '_alias/' 21 | SAVE_DIR = 'results/' + dataType_DIR + SGSModel_string + '/NX' + str(NX) + '/dt' + str(dt) + '_IC' + str(ICnum) + '/' 22 | SAVE_DIR_DATA = SAVE_DIR + 'data/' 23 | SAVE_DIR_IC = SAVE_DIR + 'IC/' 24 | 25 | # Create directories if they aren't present 26 | try: 27 | os.makedirs(SAVE_DIR_DATA) 28 | os.makedirs(SAVE_DIR_IC) 29 | except OSError as error: 30 | print(error) 31 | 32 | return SAVE_DIR, SAVE_DIR_DATA, SAVE_DIR_IC 33 | # -------------------------------------------------------------------------- 34 | def get_last_file(file_path): 35 | # Get all .mat files in the specified directory 36 | mat_files = glob.glob(os.path.join(file_path, '*.mat')) 37 | 38 | # Extract the integer values from the filenames 39 | file_numbers = [int(os.path.splitext(os.path.basename(file))[0]) for file in mat_files] 40 | 41 | # Find the highest integer value 42 | if file_numbers: 43 | last_file_number = max(file_numbers) 44 | return last_file_number 45 | else: 46 | return None 47 | return file_numbers 48 | # -------------------------------------------------------------------------- 49 | def set_last_file(last_file_number_data, last_file_number_IC): 50 | 51 | # Set last File numbers 52 | if last_file_number_data is None: 53 | print(f"Last data file number: {last_file_number_data}") 54 | last_file_number_data = 0 55 | print("Updated last data file number to " + str(last_file_number_data)) 56 | else: 57 | raise ValueError("Data already exists in the results folder for this case, either resume the simulation (resumeSim = True) or delete data to start a new simulation") 58 | 59 | if last_file_number_IC is None: 60 | print(f"Last data file number: {last_file_number_IC}") 61 | last_file_number_IC = 0 62 | print("Updated last data file number to " + str(last_file_number_IC)) 63 | else: 64 | raise ValueError("Data already exists in the results folder for this case, either resume the simulation (resumeSim = True) or delete data to start a new simulation") 65 | 66 | return last_file_number_data, last_file_number_IC 67 | # -------------------------------------------------------------------------- 68 | def save_settings(readTrue,ICnum,resumeSim,saveData, 69 | NSAVE,tSAVE,tTotal,maxit,NX, Lx, Re,dt,nu,rho,alpha, SGSModel_string, 70 | fkx, fky, SAVE_DIR): 71 | 72 | # Save variables of the solver 73 | variables = { 74 | 'readTrue': readTrue, 75 | 'ICnum': ICnum, 76 | 'resumeSim': resumeSim, 77 | 'saveData': saveData, 78 | 'NSAVE': NSAVE, 79 | 'tSAVE': tSAVE, 80 | 'tTotal': tTotal, 81 | 'maxit': maxit, 82 | 'NX': NX, 83 | 'Lx': Lx, 84 | 'Re': Re, 85 | 'dt': dt, 86 | 'nu': nu, 87 | 'rho': rho, 88 | 'alpha': alpha, 89 | 'SGSModel': SGSModel_string, 90 | 'fkx': fkx, 91 | 'fky': fky, 92 | 'SAVE_DIR': SAVE_DIR, 93 | } 94 | 95 | # Open file in write mode and save parameters 96 | with open(SAVE_DIR + 'parameters.txt', 'w') as f: 97 | # Write each variable to a new line in the file 98 | for key, value in variables.items(): 99 | f.write(f'{key}: {value}\n') 100 | 101 | print("Parameters of the flow saved to saved to " + SAVE_DIR + 'parameter.txt') 102 | 103 | def pretty_print_table(title, table_content): 104 | longest_key = max([len(str(row[0])) for row in table_content]) 105 | longest_value = max([len(str(row[1])) for row in table_content]) 106 | 107 | # Title 108 | print('+' + '-' * (longest_key+2) + '+' + '-' * (longest_value+2) + '+') 109 | print('| ' + title.center(longest_key) + ' | ' + 'Value'.center(longest_value) + ' |') 110 | print('+' + '-' * (longest_key+2) + '+' + '-' * (longest_value+2) + '+') 111 | 112 | # Content 113 | for row in table_content: 114 | print('| ' + str(row[0]).center(longest_key) + ' | ' + str(row[1]).center(longest_value) + ' |') 115 | print('+' + '-' * (longest_key+2) + '+' + '-' * (longest_value+2) + '+') 116 | print() -------------------------------------------------------------------------------- /py2d/dealias.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import jax.numpy as jnp 3 | from jax import jit 4 | 5 | from py2d.filter import coarse_spectral_filter_square, coarse_spectral_filter_square_jit 6 | from py2d.util import fft2_to_rfft2, rfft2_to_fft2 7 | 8 | def multiply_dealias(u, v, dealias=True): 9 | """ 10 | Multiply two fields in physical space, with an option to dealias the product. 11 | Reference: https://math.jhu.edu/~feilu/notes/DealiasingFFT.pdf 12 | Orszag 2/3 Rule 13 | 14 | Parameters: 15 | - u, v: 2D numpy arrays representing the fields to be multiplied. 16 | - dealias: Boolean indicating if dealiasing is required. Default is True. 17 | 18 | Returns: 19 | - Product of the two fields, optionally dealiased, in physical space. 20 | """ 21 | if dealias: 22 | 23 | Nx, Ny = u.shape 24 | 25 | # Convert to spectral space for dealiasing 26 | u_hat_alias = np.fft.rfft2(u) 27 | v_hat_alias = np.fft.rfft2(v) 28 | # dealising 29 | uv_alias_hat = multiply_dealias_spectral(u_hat_alias, v_hat_alias) 30 | uv_alias = np.fft.irfft2(uv_alias_hat, s=[Nx,Ny]) 31 | else: 32 | # Direct multiplication in physical space if no dealiasing needed 33 | uv_alias = u * v 34 | 35 | return uv_alias 36 | 37 | def multiply_dealias_spectral(u_hat, v_hat): 38 | """ 39 | Multiply two fields and dealias the result, Input is in spectral space and returned in physical space. 40 | 41 | Parameters: 42 | - u_hat, v_hat: 2D numpy arrays representing the spectral coefficients of the fields. 43 | 44 | Returns: 45 | - Product of the two fields, dealiased and converted to physical space. 46 | """ 47 | 48 | Ncoarse = u_hat.shape[0] 49 | 50 | # Dealias each field 51 | u_dealias_hat = padding_for_dealias(u_hat, spectral=True) 52 | v_dealias_hat = padding_for_dealias(v_hat, spectral=True) 53 | 54 | Nfine = u_dealias_hat.shape[0] 55 | 56 | u_dealias = np.fft.irfft2(u_dealias_hat, s=[Nfine,Nfine]) 57 | v_dealias = np.fft.irfft2(v_dealias_hat, s=[Nfine,Nfine]) 58 | 59 | u_dealias_v_dealias_hat = np.fft.rfft2(u_dealias * v_dealias) 60 | 61 | # Multiply on the dealise grid and coarse grain to alias grid 62 | uv_dealias_hat = coarse_spectral_filter_square(u_dealias_v_dealias_hat, Ncoarse) 63 | 64 | return uv_dealias_hat 65 | 66 | def padding_for_dealias(u, spectral=False, K=3/2): 67 | """ 68 | Padding zeros to the high wavenumber in the 69 | spectral space for a 2D square grid. This is required for handling 70 | nonlinear terms in simulations of 2D Homogeneous Isotropic Turbulence. 71 | Reference: https://math.jhu.edu/~feilu/notes/DealiasingFFT.pdf 72 | 73 | Parameters: 74 | - u: 2D numpy array representing the field in physical or spectral space. 75 | - spectral: Boolean indicating if 'u' is already in spectral space. Default is False. 76 | - K: Scaling factor to increase the grid size for dealiasing. Default is 3/2. 77 | 78 | Returns: 79 | - u_pad: 2D numpy array in spectral space (if spectral=True) or 80 | in physical space (if spectral=False) after padding. 81 | 82 | Note: 83 | - The u_hat_dealias needs to be conjugate symmetric for the inverse FFT to be real. 84 | - Padded dealising grid should be even in size. i.e. Ngrid_alias is multiple of 4 85 | """ 86 | # Determine the original and dealiased grid sizes 87 | 88 | if spectral: 89 | u_hat = u 90 | else: 91 | u_hat = np.fft.rfft2(u) 92 | 93 | N_coarse = u_hat.shape[0] 94 | 95 | if N_coarse % 2 == 0: 96 | N_pad = int(K * N_coarse) 97 | else: 98 | N_pad = int(K * (N_coarse-1)) + 1 99 | 100 | 101 | # Index of Nyquist wave number - fine grid 102 | kn_pad = N_pad//2 103 | # Index of Nyquist wave number - coarse grid 104 | kn_coarse =int(N_coarse//2) 105 | 106 | # Scale the spectral data to account for the increased grid size 107 | u_hat_scaled = (N_pad/N_coarse)**2 * u_hat 108 | 109 | # Initialize a 2D array of zeros with the dealiased shape 110 | u_hat_pad = np.zeros((N_pad, kn_pad+1), dtype=complex) 111 | 112 | # Compute indices for padding 113 | if N_coarse % 2 == 0: 114 | indx_pad = np.r_[0:kn_coarse+1, N_pad-kn_coarse+1:N_pad] 115 | else: 116 | indx_pad = np.r_[0:kn_coarse+1, N_pad-kn_coarse:N_pad] 117 | indy_pad = np.r_[0:kn_coarse+1] 118 | 119 | # Apply scaling and pad the spectral data with zeros 120 | u_hat_pad[np.ix_(indx_pad, indy_pad)] = u_hat_scaled 121 | 122 | # ################## Making the data conjugate symmetric ################## 123 | 124 | if N_pad % 2 == 0: 125 | 126 | # # ** Method#1: Making the array conjugate symmetric 127 | # u_hat_pad[N_pad-N_coarse//2,:] = np.conj(u_hat_pad[N_coarse//2,:]) 128 | 129 | # # ** Method #2: Making the Nyquist wavenumber 0 130 | u_hat_pad[kn_coarse,:] = 0 131 | u_hat_pad[:,kn_coarse] = 0 132 | pass 133 | 134 | else: # Odd grid size 135 | 136 | # ** Method#1: Do nothing - its already conjugate symmetric 137 | 138 | # # # ** Method #2: Making the Nyquist wavenumber 0 139 | u_hat_pad[kn_coarse,:] = 0 140 | u_hat_pad[kn_pad+(kn_pad-kn_coarse)+1,:] = 0 141 | u_hat_pad[:,kn_coarse] = 0 142 | pass 143 | 144 | # ** Method#3: take irfft and back: works for both odd and even grid sizes 145 | # u = np.fft.irfft2(u_hat_pad, s=[N_dealias,N_dealias]) 146 | # u_hat_pad = np.fft.rfft2(u) 147 | 148 | # return u_hat_dealias 149 | if spectral: 150 | return u_hat_pad 151 | else: 152 | return np.fft.irfft2(u_hat_pad, s=[N_pad,N_pad]) 153 | 154 | # def padding_for_dealias(u, spectral=False, K=3/2): 155 | # """ 156 | # Padding zeros to the high wavenumber in the 157 | # spectral space for a 2D square grid. This is required for handling 158 | # nonlinear terms in simulations of 2D Homogeneous Isotropic Turbulence. 159 | # Reference: https://math.jhu.edu/~feilu/notes/DealiasingFFT.pdf 160 | 161 | # Parameters: 162 | # - u: 2D numpy array representing the field in physical or spectral space. 163 | # - spectral: Boolean indicating if 'u' is already in spectral space. Default is False. 164 | # - K: Scaling factor to increase the grid size for dealiasing. Default is 3/2. 165 | 166 | # Returns: 167 | # - u_dealias: 2D numpy array in spectral space (if spectral=True) or 168 | # in physical space (if spectral=False) after padding. 169 | 170 | # Note: 171 | # - The u_hat_dealias needs to be conjugate symmetric for the inverse FFT to be real. 172 | # - Padded dealising grid should be even in size. i.e. Ngrid_alias is multiple of 4 173 | # """ 174 | # # Determine the original and dealiased grid sizes 175 | # N_alias = u.shape[0] 176 | # N_dealias = int(K * N_alias) 177 | 178 | # if np.mod(N_dealias,2) != 0: 179 | # raise ValueError("Padded dealising grid should be even in size. Dealiased grid size is ", N_dealias) 180 | 181 | # if spectral: 182 | # u_hat_alias = u 183 | # else: 184 | # # Compute the spectral coefficients of the input field 185 | # u_hat_alias = np.fft.fft2(u) 186 | 187 | # # Scale the spectral data to account for the increased grid size 188 | # u_hat_alias_scaled = K**2 * u_hat_alias 189 | 190 | # # Initialize a 2D array of zeros with the dealiased shape 191 | # u_hat_dealias = np.zeros((N_dealias, N_dealias), dtype=complex) 192 | 193 | # # Compute indices for padding 194 | # indvpad = np.r_[0:int(N_alias/2)+1, N_dealias-int(N_alias/2)+1:N_dealias] 195 | 196 | # # Apply scaling and pad the spectral data with zeros 197 | # u_hat_dealias[np.ix_(indvpad, indvpad)] = u_hat_alias_scaled 198 | 199 | # # Making the padded array conjugate symmetric 200 | # u_hat_dealias_pad = conjugate_symmetrize_padding(u_hat_dealias, K=K) 201 | 202 | # if spectral: 203 | # return u_hat_dealias_pad 204 | # else: 205 | # # Return in the appropriate space 206 | # return np.fft.ifft2(u_hat_dealias_pad).real 207 | 208 | # def conjugate_symmetrize_padding(a_hat_fft_pad, K): 209 | 210 | # a_hat_fft_pad_sym = a_hat_fft_pad.copy() 211 | 212 | # # Nyquist wavenumber index after Padding - This is the Nyquist wavenumber of the dealiased grid 213 | # Ny = int(a_hat_fft_pad.shape[0])//2 214 | 215 | # # Nyquist wavenumber index before Padding - This is the Nyquist wavenumber of the original grid 216 | # Ny_old = int(Ny*1/K) 217 | 218 | # # ######################## Method #1 ############################ 219 | # # # 0th Wavenumber is conjugate symmetric 220 | # # a_hat_fft_pad_sym[0,Ny+1:] = np.flip(a_hat_fft_pad[0,1:Ny]).conj() 221 | # # a_hat_fft_pad_sym[Ny+1:,0] = np.flip(a_hat_fft_pad[1:Ny,0]).conj() 222 | 223 | # # # Nyquist wavenumber - before padding (Ny_old) is conjugate symmetric 224 | # # # Padding creates this un-symmetricity 225 | # # a_hat_fft_pad_sym[Ny+(Ny-Ny_old),1:] = np.flip(a_hat_fft_pad[Ny-(Ny-Ny_old),1:]).conj() 226 | # # a_hat_fft_pad_sym[1:,Ny+(Ny-Ny_old)] = np.flip(a_hat_fft_pad[1:,Ny-(Ny-Ny_old)]).conj() 227 | 228 | # # # 0th wavenumber 0th wavenumbers has zero imaginary part 229 | # # a_hat_fft_pad_sym[0,0] = a_hat_fft_pad[0,0].real # (Kx=0, Ky=0) 230 | 231 | # # # Nyquist wavenumber of 0th wavenumber has zero imaginary part 232 | # # a_hat_fft_pad_sym[0,Ny//2] = a_hat_fft_pad[0,Ny//2].real # (Kx=0, Ky=Ny) 233 | # # a_hat_fft_pad_sym[Ny//2,0] = a_hat_fft_pad[Ny//2,0].real # (Kx=Ny, Ky=0) 234 | 235 | # # # Nyquist wavenumber of Nyquist wavenumber has zero imaginary part 236 | # # a_hat_fft_pad_sym[Ny,Ny] = a_hat_fft_pad[Ny,Ny].real # (Kx=Ny, Ky=Ny) 237 | 238 | # ####################### Method #2 ############################ 239 | 240 | # # ############ Alternatively equate the data to zero at Ny_old wavenumber (Kx,Ky) = (Ny-(Ny-Ny_old),Ny+(Ny-Ny_old)) to make it conjugate symmetric 241 | # a_hat_fft_pad_sym[Ny_old,:] = 0 242 | # a_hat_fft_pad_sym[:,Ny_old] = 0 243 | 244 | # # # This is already zero - don't need to set it again 245 | # # a_hat_fft_pad_sym[Ny+(Ny-Ny_old),:] = 0 246 | # # a_hat_fft_pad_sym[:,Ny+(Ny-Ny_old)] = 0 247 | 248 | # return a_hat_fft_pad_sym 249 | 250 | @jit 251 | def multiply_dealias_physical_jit(a, b): 252 | """ 253 | Multiply two fields and dealias the result, Input is in spectral space and returned in physical space. 254 | 255 | Parameters: 256 | - u_hat, v_hat: 2D numpy arrays representing the spectral coefficients of the fields. 257 | 258 | Returns: 259 | - Product of the two fields, dealiased and converted to physical space. 260 | """ 261 | Nx, Ny = a.shape 262 | 263 | a_hat = jnp.fft.rfft2(a) 264 | b_hat = jnp.fft.rfft2(b) 265 | 266 | ab_dealias_hat = multiply_dealias_spectral_jit(a_hat, b_hat) 267 | 268 | ab_dealias = jnp.fft.irfft2(ab_dealias_hat, s=[Nx,Ny]) 269 | 270 | return ab_dealias 271 | 272 | @jit 273 | def multiply_dealias_spectral_jit(a_hat, b_hat): 274 | """ 275 | Multiply two fields and dealias the result, Input is in spectral space and returned in physical space. 276 | 277 | Parameters: 278 | - u_hat, v_hat: 2D numpy arrays representing the spectral coefficients of the fields. 279 | 280 | Returns: 281 | - Product of the two fields, dealiased and converted to physical space. 282 | """ 283 | 284 | Ncoarse = a_hat.shape[0] 285 | 286 | # Dealias each field 287 | a_dealias_hat = padding_for_dealias_spectral_jit(a_hat) 288 | b_dealias_hat = padding_for_dealias_spectral_jit(b_hat) 289 | 290 | Nfine = a_dealias_hat.shape[0] 291 | 292 | a_dealias = jnp.fft.irfft2(a_dealias_hat, s=[Nfine,Nfine]) 293 | b_dealias = jnp.fft.irfft2(b_dealias_hat, s=[Nfine,Nfine]) 294 | 295 | a_dealias_b_dealias_hat = jnp.fft.rfft2(a_dealias * b_dealias) 296 | 297 | # Multiply on the dealise grid and coarse grain to alias grid 298 | ab_dealias_hat = coarse_spectral_filter_square_jit(a_dealias_b_dealias_hat, Ncoarse) 299 | 300 | return ab_dealias_hat 301 | 302 | 303 | def padding_for_dealias_spectral_jit(u_hat, K=3/2): 304 | """ 305 | Pads zeros to the high wavenumbers in spectral space for dealiasing in 2D 306 | simulations. 307 | 308 | This function is designed to be JAX-compatible and includes multiple methods 309 | to ensure conjugate symmetry for real-valued iFFT results. 310 | 311 | Args: 312 | u_hat: A JAX array representing the field in spectral space. 313 | spectral: Boolean indicating if 'u' is already in spectral space. 314 | Default is False. 315 | K: Scaling factor to increase the grid size for dealiasing. Default is 3/2. 316 | 317 | Returns: 318 | u_hat_pad: A JAX array in spectral space (if spectral=True) or in physical 319 | space (if spectral=False) after padding. 320 | """ 321 | 322 | N_coarse = u_hat.shape[0] 323 | N_pad = int(K * N_coarse) 324 | 325 | kn_pad = N_pad // 2 326 | kn_coarse = N_coarse // 2 327 | 328 | u_hat_scaled = (N_pad / N_coarse) ** 2 * u_hat 329 | u_hat_pad = jnp.zeros((N_pad, kn_pad + 1), dtype=complex) 330 | 331 | 332 | u_hat_pad = u_hat_pad.at[:kn_coarse+1, :kn_coarse+1].set(u_hat_scaled[:kn_coarse+1, :kn_coarse+1]) 333 | 334 | if N_pad % 2 == 0: 335 | u_hat_pad = u_hat_pad.at[N_pad-kn_coarse+1:,:kn_coarse+1].set(u_hat_scaled[N_coarse-kn_coarse+1:, :kn_coarse+1]) 336 | 337 | else: 338 | u_hat_pad = u_hat_pad.at[N_pad-kn_coarse:,:kn_coarse+1].set(u_hat_scaled[N_coarse-kn_coarse:, :kn_coarse+1]) 339 | 340 | # ################## Making the data conjugate symmetric ################## 341 | 342 | if N_pad % 2 == 0: 343 | 344 | # # Method #1: Direct manipulation 345 | # u_hat_pad = u_hat_pad.at[N_pad - N_coarse // 2, :].set(jnp.conj(u_hat_pad[N_coarse // 2, :])) 346 | 347 | # # # Method #2: Setting Nyquist wavenumbers to zero 348 | u_hat_pad = u_hat_pad.at[kn_coarse, :].set(0) 349 | u_hat_pad = u_hat_pad.at[:, kn_coarse].set(0) 350 | 351 | else: # Odd grid size 352 | 353 | # ** Method#1: Do nothing - its already conjugate symmetric 354 | 355 | # # # Method #2: Setting Nyquist wavenumbers to zero 356 | u_hat_pad = u_hat_pad.at[kn_coarse, :].set(0) 357 | u_hat_pad = u_hat_pad.at[kn_pad+(kn_pad-kn_coarse)+1, :].set(0) 358 | u_hat_pad = u_hat_pad.at[:, kn_coarse].set(0) 359 | pass 360 | 361 | # # # Method #3: iFFT and FFT roundtrip 362 | # # u = jnp.fft.irfft2(u_hat_pad, s=[N_pad, N_pad]) 363 | # # u_hat_pad = jnp.fft.rfft2(u) 364 | 365 | return u_hat_pad 366 | 367 | # @jit 368 | # def padding_for_dealias_spectral_jit(a_hat_alias, K=3/2): 369 | # """ 370 | # Padding zeros to the high wavenumber in the 371 | # spectral space for a 2D square grid. This is required for handling 372 | # nonlinear terms in simulations of 2D Homogeneous Isotropic Turbulence. 373 | # Reference: https://math.jhu.edu/~feilu/notes/DealiasingFFT.pdf 374 | 375 | # Parameters: 376 | # - u: 2D numpy array representing the field in spectral space. 377 | # - spectral: Boolean indicating if 'u' is already in spectral space. Default is False. 378 | # - K: Scaling factor to increase the grid size for dealiasing. Default is 3/2. 379 | 380 | # Returns: 381 | # - u_dealias_hat: 2D numpy array in spectral space after padding. 382 | 383 | # Note: 384 | # - The u_hat_dealias needs to be conjugate symmetric for the inverse FFT to be real. 385 | # """ 386 | # # Determine the original and dealiased grid sizes 387 | # N_alias = a_hat_alias.shape[0] 388 | # N_dealias = int(K * N_alias) 389 | 390 | # # Scale the spectral data to account for the increased grid size 391 | # a_hat_alias_scaled = K**2 * a_hat_alias 392 | 393 | # # ********** Jax Code ************ 394 | 395 | # a_hat_dealias = jnp.zeros((N_dealias, N_dealias), dtype=complex) 396 | 397 | # # Extract the corners of the scaled array 398 | # utopleft = a_hat_alias_scaled[0:int(N_alias/2)+1, 0:int(N_alias/2)+1] 399 | # utopright = a_hat_alias_scaled[0:int(N_alias/2)+1, N_alias-int(N_alias/2)+1:N_alias] 400 | # ubottomleft = a_hat_alias_scaled[N_alias-int(N_alias/2)+1:N_alias, 0:int(N_alias/2)+1] 401 | # ubottomright = a_hat_alias_scaled[N_alias-int(N_alias/2)+1:N_alias, N_alias-int(N_alias/2)+1:N_alias] 402 | 403 | # # Since JAX arrays are immutable, use the .at[].set() method for updates 404 | # a_hat_dealias = a_hat_dealias.at[0:int(N_alias/2)+1, 0:int(N_alias/2)+1].set(utopleft) 405 | # a_hat_dealias = a_hat_dealias.at[0:int(N_alias/2)+1, N_dealias-int(N_alias/2)+1:N_dealias].set(utopright) 406 | # a_hat_dealias = a_hat_dealias.at[N_dealias-int(N_alias/2)+1:N_dealias, 0:int(N_alias/2)+1].set(ubottomleft) 407 | # a_hat_dealias = a_hat_dealias.at[N_dealias-int(N_alias/2)+1:N_dealias, N_dealias-int(N_alias/2)+1:N_dealias].set(ubottomright) 408 | 409 | # ######################## Making array conjugate symmetric ############################ 410 | # # Making first row,col conjugate symmetric as well as the Nyquist +/- 1 row,col 411 | 412 | # a_hat_fft_pad_sym = a_hat_dealias.copy() 413 | 414 | # # Nyquist wavenumber index after Padding - This is the Nyquist wavenumber of the dealiased grid 415 | # Ny = int(a_hat_dealias.shape[0])//2 416 | 417 | # # Nyquist wavenumber index before Padding - This is the Nyquist wavenumber of the original grid 418 | # Ny_old = int(Ny*1/K) 419 | 420 | # # First Row and Column (using at.set) 421 | # # a_hat_fft_pad_sym = a_hat_fft_pad_sym.at[0, N_dealias//2+1].set(jnp.conj(a_hat_fft_pad_sym[0, N_dealias//2-1])) 422 | # # a_hat_fft_pad_sym = a_hat_fft_pad_sym.at[N_dealias//2+1, 0].set(jnp.conj(a_hat_fft_pad_sym[N_dealias//2-1, 0])) 423 | 424 | # # # Conjugate Symmetry with Slicing and Vectorization 425 | # # # Calculate slice indices dynamically 426 | # # slice_start = N_dealias // 2 - 1 427 | # # slice_end = -slice_start 428 | 429 | # # a_hat_fft_pad_sym = a_hat_fft_pad_sym.at[slice_end, 1:].set(jnp.conj(jnp.flip(a_hat_fft_pad_sym[slice_start, 1:]))) 430 | # # a_hat_fft_pad_sym = a_hat_fft_pad_sym.at[1:, slice_end].set(jnp.conj(jnp.flip(a_hat_fft_pad_sym[1:, slice_start]))) 431 | 432 | # # # ############ Alternatively equate the data to zero at wavenumber (N//2-1) to (N//2+1) to make it conjugate symmetric 433 | # a_hat_fft_pad_sym = a_hat_fft_pad_sym.at[Ny_old,:].set(0) 434 | # a_hat_fft_pad_sym = a_hat_fft_pad_sym.at[:,Ny_old].set(0) 435 | 436 | # # # This is already zero - don't need to set it again 437 | # # a_hat_fft_pad_sym = a_hat_fft_pad_sym.at[Ny+(Ny-Ny_old),:].set(0) 438 | # # a_hat_fft_pad_sym = a_hat_fft_pad_sym.at[:,Ny+(Ny-Ny_old)].set(0) 439 | 440 | # # Return in the appropriate space 441 | # return a_hat_fft_pad_sym 442 | 443 | -------------------------------------------------------------------------------- /py2d/derivative.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def derivative(T, order, Kx, Ky, spectral=False): 4 | """ 5 | Calculate spatial derivatives for 2D_FHIT in spectral space. 6 | Boundary conditions are periodic in x and y spatial dimensions 7 | 8 | Parameters: 9 | ---------- 10 | T : numpy.ndarray 11 | Input flow field. If `spectral` is False, T is in the physical domain; otherwise, it is in the spectral domain. 12 | order : list 13 | Array of order of derivatives in x and y spatial dimensions. Expects a list of two integers >=0. 14 | Kx, Ky : numpy.ndarray 15 | Precomputed wavenumbers in the x and y dimensions. 16 | spectral : bool, optional 17 | Whether T is in the spectral domain. Default is False. 18 | 19 | Returns: 20 | ------- 21 | Tderivative or Tderivative_hat : numpy.ndarray 22 | The derivative of the flow field T. If `spectral` is False, this is in the physical domain; otherwise, it's in the spectral domain. 23 | """ 24 | 25 | # If the input data is not in the spectral domain, we transform it 26 | if spectral == False: 27 | Tderivative = derivative_physical(T, order, Kx, Ky) 28 | return Tderivative 29 | else: # if it is already in the spectral domain, we do nothing 30 | T_hat = T 31 | Tderivative_hat = derivative_spectral(T_hat, order, Kx, Ky) 32 | return Tderivative_hat 33 | 34 | 35 | def derivative_spectral(T_hat, order, Kx, Ky): 36 | """ 37 | Calculate spatial derivatives for 2D_FHIT in spectral space. 38 | Boundary conditions are periodic in x and y spatial dimensions 39 | 40 | Parameters: 41 | ---------- 42 | T_hat : numpy.ndarray 43 | Input flow field. T_hat is in the spectral domain. 44 | order : list 45 | Array of order of derivatives in x and y spatial dimensions. Expects a list of two integers >=0. 46 | Kx, Ky : numpy.ndarray 47 | Precomputed wavenumbers in the x and y dimensions. 48 | 49 | Returns: 50 | ------- 51 | Tderivative_hat : numpy.ndarray 52 | The derivative of the flow field T_hat in the spectral domain. 53 | """ 54 | 55 | # Orders of derivatives in x and y dimensions 56 | orderX = order[0] 57 | orderY = order[1] 58 | 59 | # Calculating derivatives in spectral space using the Fourier derivative theorem 60 | Tderivative_hat = ((1j*Kx)**orderX) * ((1j*Ky)**orderY) * T_hat 61 | 62 | return Tderivative_hat 63 | 64 | def derivative_physical(T, order, Kx, Ky, ): 65 | """ 66 | Calculate spatial derivatives for 2D_FHIT in physical space using spectral methods. 67 | Boundary conditions are periodic in x and y spatial dimensions 68 | 69 | Parameters: 70 | ---------- 71 | T : numpy.ndarray 72 | Input flow field. T is in the physical domain. 73 | order : list 74 | Array of order of derivatives in x and y spatial dimensions. Expects a list of two integers >=0. 75 | Kx, Ky : numpy.ndarray 76 | Precomputed wavenumbers in the x and y dimensions. 77 | 78 | Returns: 79 | ------- 80 | Tderivative : numpy.ndarray 81 | The derivative of the flow field T in the physical domain. 82 | """ 83 | 84 | Nx, Ny = T.shape 85 | 86 | # We transform data to spectral domain 87 | T_hat = np.fft.rfft2(T) 88 | 89 | # Orders of derivatives in x and y dimensions 90 | Tderivative_hat = derivative_spectral(T_hat, order, Kx, Ky) 91 | 92 | # Transform the result back into the physical domain 93 | Tderivative = np.fft.irfft2(Tderivative_hat, s=[Nx, Ny]) 94 | 95 | 96 | 97 | return Tderivative 98 | 99 | 100 | 101 | 102 | -------------------------------------------------------------------------------- /py2d/eddy_viscosity_models.py: -------------------------------------------------------------------------------- 1 | # ---------------------------------------------------------------------- 2 | # Created : Karan Jakhar May 2023 3 | # ---------------------------------------------------------------------- 4 | 5 | # Eddy Viscosity SGS Models for 2D Turbulence solver 6 | 7 | 8 | import numpy as np 9 | import jax.numpy as jnp 10 | from jax import jit 11 | 12 | from py2d.convert import strain_rate_spectral 13 | from py2d.filter import spectral_filter_square_same_size_jit 14 | 15 | strain_rate_spectral_jit = jit(strain_rate_spectral) 16 | 17 | def Tau_eddy_viscosity(eddy_viscosity, Psi_hat, Kx, Ky): 18 | ''' 19 | Calculate the eddy viscosity term (Tau) in the momentum equation 20 | ''' 21 | N = Psi_hat.shape[0] 22 | 23 | S11_hat, S12_hat, S22_hat = strain_rate_spectral(Psi_hat, Kx, Ky) 24 | S11 = np.fft.irfft2(S11_hat, s=[N, N]) 25 | S12 = np.fft.irfft2(S12_hat, s=[N, N]) 26 | S22 = np.fft.irfft2(S22_hat, s=[N, N]) 27 | 28 | Tau11 = -2*eddy_viscosity*S11 29 | Tau12 = -2*eddy_viscosity*S12 30 | Tau22 = -2*eddy_viscosity*S22 31 | 32 | return Tau11, Tau12, Tau22 33 | 34 | def Sigma_eddy_viscosity(eddy_viscosity, Omega_hat, Kx, Ky): 35 | ''' 36 | Calculate the eddy viscosity term (Tau) in the momentum equation 37 | ''' 38 | N = Omega_hat.shape[0] 39 | 40 | Omegax_hat = (1.j) * Kx * Omega_hat 41 | Omegay_hat = (1.j) * Ky * Omega_hat 42 | 43 | Omegax = np.fft.irfft2(Omegax_hat, s=[N, N]) 44 | Omegay = np.fft.irfft2(Omegay_hat, s=[N, N]) 45 | 46 | Sigma1 = -eddy_viscosity*Omegax 47 | Sigma2 = -eddy_viscosity*Omegay 48 | 49 | return Sigma1, Sigma2 50 | 51 | @jit 52 | def eddy_viscosity_smag(Cs, Delta, characteristic_S): 53 | ''' 54 | Smagorinsky Model (SMAG) 55 | ''' 56 | characteristic_S2 = characteristic_S ** 2 57 | characteristic_S_mean = jnp.sqrt(jnp.mean(characteristic_S2)) 58 | eddy_viscosity = (Cs * Delta) ** 2 * characteristic_S_mean 59 | return eddy_viscosity 60 | 61 | @jit 62 | def eddy_viscosity_smag_local(Cs, Delta, characteristic_S): 63 | ''' 64 | Smagorinsky Model (SMAG) - Local characteristic_S 65 | ''' 66 | characteristic_S2 = characteristic_S ** 2 67 | characteristic_S = jnp.sqrt(characteristic_S2) 68 | eddy_viscosity = (Cs * Delta) ** 2 * characteristic_S 69 | return eddy_viscosity 70 | 71 | @jit 72 | def characteristic_strain_rate_smag(Psi_hat, Kx, Ky, Ksq): 73 | ''' 74 | Characteristic strain rate 75 | Required for Smagorinsky Model 76 | ''' 77 | N = Psi_hat.shape[0] 78 | 79 | (S11_hat, S12_hat, _) = strain_rate_spectral_jit(Psi_hat, Kx, Ky) 80 | S11 = jnp.fft.irfft2(S11_hat, s=[N, N]) 81 | S12 = jnp.fft.irfft2(S12_hat, s=[N, N]) 82 | characteristic_S = 2 * jnp.sqrt(S11 ** 2 + S12 ** 2) 83 | return characteristic_S 84 | 85 | @jit 86 | def eddy_viscosity_leith(Cl, Delta, characteristic_Omega): 87 | ''' 88 | Leith Model (LEITH) 89 | ''' 90 | characteristic_Omega_mean = jnp.mean(characteristic_Omega) 91 | ls = Cl * Delta 92 | eddy_viscosity = ls ** 3 * characteristic_Omega_mean 93 | return eddy_viscosity 94 | 95 | @jit 96 | def characteristic_omega_leith(Omega_hat, Kx, Ky): 97 | ''' 98 | Characteristic Gradient of Omega 99 | Required for Leith Model 100 | ''' 101 | N = Omega_hat.shape[0] 102 | 103 | Omegax_hat = (1.j) * Kx * Omega_hat 104 | Omegay_hat = (1.j) * Ky * Omega_hat 105 | Omegax = jnp.fft.irfft2(Omegax_hat, s=[N, N]) 106 | Omegay = jnp.fft.irfft2(Omegay_hat, s=[N, N]) 107 | characteristic_Omega = jnp.sqrt(Omegax ** 2 + Omegay ** 2) 108 | return characteristic_Omega 109 | 110 | @jit 111 | def coefficient_dsmag_PsiOmega(Psi_hat, Omega_hat, characteristic_S, Kx, Ky, Ksq, Delta): 112 | ''' 113 | cs = Cs**2 114 | Dynamic Coefficient for Dynamic Smagorinsky Model (DSMAG) 115 | ''' 116 | (Psif_hat, Omegaf_hat, Omega_lap, Omegaf_lap, Delta_test, nx_test) = initialize_filtered_variables_PsiOmega( 117 | Psi_hat, Omega_hat, Ksq, Delta) 118 | L = residual_jacobian_PsiOmega(Psi_hat, Omega_hat, Psif_hat, Omegaf_hat, Kx, Ky, nx_test) 119 | M = residual_dsmag_PsiOmega(Omega_lap, Omegaf_lap, characteristic_S, Delta, Delta_test, nx_test) 120 | cs = coefficient_dynamic_PsiOmega(L, M) 121 | return cs 122 | 123 | @jit 124 | def coefficient_dsmaglocal_PsiOmega(Psi_hat, Omega_hat, characteristic_S, Kx, Ky, Ksq, Delta): 125 | ''' 126 | cs = Cs**2 127 | Dynamic Coefficient for Dynamic Smagorinsky Model (DSMAG) with local Cs 128 | ''' 129 | (Psif_hat, Omegaf_hat, Omega_lap, Omegaf_lap, Delta_test, nx_test) = initialize_filtered_variables_PsiOmega( 130 | Psi_hat, Omega_hat, Ksq, Delta) 131 | L = residual_jacobian_PsiOmega(Psi_hat, Omega_hat, Psif_hat, Omegaf_hat, Kx, Ky, nx_test) 132 | M = residual_dsmag_PsiOmega(Omega_lap, Omegaf_lap, characteristic_S, Delta, Delta_test, nx_test) 133 | cs = coefficient_dynamiclocal_PsiOmega(L, M) 134 | return cs 135 | 136 | @jit 137 | def coefficient_dleith_PsiOmega(Psi_hat, Omega_hat, characteristic_Omega, Kx, Ky, Ksq, Delta): 138 | ''' 139 | cl = Cl**3 140 | Dynamic Coefficient for Dynamic Leith Model (DLEITH) 141 | ''' 142 | (Psif_hat, Omegaf_hat, Omega_lap, Omegaf_lap, Delta_test, nx_test) = initialize_filtered_variables_PsiOmega( 143 | Psi_hat, Omega_hat, Ksq, Delta) 144 | L = residual_jacobian_PsiOmega(Psi_hat, Omega_hat, Psif_hat, Omegaf_hat, Kx, Ky, nx_test) 145 | M = residual_dleith_PsiOmega(Omega_lap, Omegaf_lap, characteristic_Omega, Delta, Delta_test, nx_test) 146 | cl = coefficient_dynamic_PsiOmega(L, M) 147 | return cl 148 | 149 | @jit 150 | def coefficient_dleithlocal_PsiOmega(Psi_hat, Omega_hat, characteristic_Omega, Kx, Ky, Ksq, Delta): 151 | ''' 152 | cl = Cl**3 153 | Dynamic Coefficient for Dynamic Leith Model (DLEITH) 154 | ''' 155 | (Psif_hat, Omegaf_hat, Omega_lap, Omegaf_lap, Delta_test, nx_test) = initialize_filtered_variables_PsiOmega( 156 | Psi_hat, Omega_hat, Ksq, Delta) 157 | L = residual_jacobian_PsiOmega(Psi_hat, Omega_hat, Psif_hat, Omegaf_hat, Kx, Ky, nx_test) 158 | M = residual_dleith_PsiOmega(Omega_lap, Omegaf_lap, characteristic_Omega, Delta, Delta_test, nx_test) 159 | cl = coefficient_dynamiclocal_PsiOmega(L, M) 160 | return cl 161 | 162 | @jit 163 | def coefficient_dynamic_PsiOmega(L, M): 164 | ''' 165 | Dynamic Coefficient 166 | Required for DSMAG and DLEITH model 167 | ''' 168 | LM = L * M 169 | MM = M * M 170 | LM_pos = 0.5 * (LM + jnp.abs(LM)) 171 | 172 | c_dynamic = jnp.mean(LM_pos) / jnp.mean(MM) 173 | return c_dynamic 174 | 175 | 176 | @jit 177 | def coefficient_dynamiclocal_PsiOmega(L, M): 178 | ''' 179 | Dynamic Coefficient 180 | Required for DSMAG and DLEITH model 181 | Attemps for: LOCAL 182 | ''' 183 | LM = L * M 184 | MM = M * M 185 | LM_pos = 0.5 * (LM + jnp.abs(LM)) 186 | 187 | #c_dynamic = np.mean(LM_pos) / np.mean(MM) 188 | c_dynamic = (LM_pos) / jnp.mean(MM) 189 | return c_dynamic 190 | 191 | @jit 192 | def initialize_filtered_variables_PsiOmega(Psi_hat, Omega_hat, Ksq, Delta): 193 | ''' 194 | Calculate filtered flow variables for DSMAG and DLEITH model - Using Psi and Omega 195 | ''' 196 | nx = Psi_hat.shape[0] 197 | nx_test = nx // 2 198 | Delta_test = 2 * Delta 199 | Psif_hat = spectral_filter_square_same_size_jit(Psi_hat, kc=nx_test//2) 200 | Omegaf_hat = spectral_filter_square_same_size_jit(Omega_hat, kc=nx_test//2) 201 | Omega_lap = jnp.fft.irfft2(-Ksq * Omega_hat, s=[nx,nx]) 202 | Omegaf_lap = jnp.fft.irfft2(-Ksq * Omegaf_hat, s=[nx,nx]) 203 | return Psif_hat, Omegaf_hat, Omega_lap, Omegaf_lap, Delta_test, nx_test 204 | 205 | @jit 206 | def residual_jacobian_PsiOmega(Psi_hat, Omega_hat, Psif_hat, Omegaf_hat, Kx, Ky, nx_test): 207 | ''' 208 | Residual of Jacobian 209 | Difference between Filtered Jacobian and Jacobian of Filtered flow variables 210 | ''' 211 | N = Psi_hat.shape[0] 212 | 213 | J1 = jacobian_Spectral2Physical(Omega_hat, Psi_hat, Kx, Ky) 214 | J1f_hat = spectral_filter_square_same_size_jit(jnp.fft.rfft2(J1), kc=nx_test//2) 215 | J1f = jnp.fft.irfft2(J1f_hat, s=[N,N]) 216 | J2f = jacobian_Spectral2Physical(Omegaf_hat, Psif_hat, Kx, Ky) 217 | L = J1f - J2f 218 | return L 219 | 220 | @jit 221 | def residual_dsmag_PsiOmega(Omega_lap, Omegaf_lap, characteristic_S, Delta, Delta_test, nx_test): 222 | ''' 223 | Residual of SMAG 224 | Difference between Filtered SMAG and SMAG of Filtered flow variables 225 | Required for DSMAG 226 | ''' 227 | N = Omega_lap.shape[0] 228 | 229 | M1_hat = spectral_filter_square_same_size_jit(jnp.fft.rfft2(characteristic_S * Omega_lap), kc=nx_test//2) 230 | M1 = Delta ** 2 * jnp.fft.irfft2(M1_hat, s=[N,N]) 231 | characteristic_Sf_hat = spectral_filter_square_same_size_jit(jnp.fft.rfft2(characteristic_S), kc=nx_test//2) 232 | characteristic_Sf = jnp.fft.irfft2(characteristic_Sf_hat, s=[N,N]) 233 | M2 = Delta_test ** 2 * characteristic_Sf * Omegaf_lap 234 | M = M1 - M2 235 | return M 236 | 237 | @jit 238 | def residual_dleith_PsiOmega(Omega_lap, Omegaf_lap, characteristic_Omega, Delta, Delta_test, nx_test): 239 | ''' 240 | Residual of LEITH 241 | Difference between Filtered LEITH and LEITH of Filtered flow variables 242 | Required for DLEITH 243 | ''' 244 | N = Omega_lap.shape[0] 245 | 246 | M1_hat = spectral_filter_square_same_size_jit(jnp.fft.rfft2(characteristic_Omega * Omega_lap), kc=nx_test//2) 247 | M1 = Delta ** 3 * jnp.fft.irfft2(M1_hat, s=[N,N]) 248 | characteristic_Omegaf_hat = spectral_filter_square_same_size_jit(jnp.fft.rfft2(characteristic_Omega), kc=nx_test//2) 249 | characteristic_Omegaf = jnp.fft.irfft2(characteristic_Omegaf_hat, s=[N,N]) 250 | M2 = Delta_test ** 3 * characteristic_Omegaf * Omegaf_lap 251 | M = M1 - M2 252 | return M 253 | 254 | @jit 255 | def jacobian_Spectral2Physical(a_hat, b_hat, Kx, Ky): 256 | """ 257 | Calculate the Jacobian (in physical space) of two scalar fields (spectral space) a and b for 2DFHIT. 258 | 259 | Parameters: 260 | ----------- 261 | a_hat : numpy.ndarray 262 | First scalar field (2D array) in spectral space, depending on the 'spectral' flag. 263 | b_hat : numpy.ndarray 264 | Second scalar field (2D array) in spectral space, depending on the 'spectral' flag. 265 | Kx : numpy.ndarray 266 | 2D array of wavenumbers in the x-direction. 267 | Ky : numpy.ndarray 268 | 2D array of wavenumbers in the y-direction. 269 | 270 | Returns: 271 | -------- 272 | J : numpy.ndarray 273 | Jacobian of scalar fields a and b (2D array) in physical space. 274 | """ 275 | N = a_hat.shape[0] 276 | 277 | ax_hat = (1.j)*Kx * a_hat 278 | ay_hat = (1.j)*Ky * a_hat 279 | bx_hat = (1.j)*Kx * b_hat 280 | by_hat = (1.j)*Ky * b_hat 281 | ax = jnp.fft.irfft2(ax_hat, s=[N,N]) 282 | ay = jnp.fft.irfft2(ay_hat, s=[N,N]) 283 | bx = jnp.fft.irfft2(bx_hat, s=[N,N]) 284 | by = jnp.fft.irfft2(by_hat, s=[N,N]) 285 | J = ax * by - ay * bx 286 | return J -------------------------------------------------------------------------------- /py2d/filter.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import jax.numpy as jnp 3 | from jax import jit 4 | 5 | from py2d.initialize import initialize_wavenumbers_rfft2 6 | from py2d.util import fft2_to_rfft2, rfft2_to_fft2 7 | 8 | def filter2D(U, filterType='gaussian', coarseGrainType='spectral', Delta=None, Ngrid=None, spectral=False): 9 | """ 10 | Filters and coarse grains 2D square grids typically used in 2D Forced Homogeneous Isotropic Turbulence (2D-FHIT). 11 | 12 | Parameters 13 | ---------- 14 | U : numpy.ndarray 15 | The 2D input data to be filtered and coarse grained. 16 | 17 | filterType : str, optional 18 | The type of filter to apply. It can be 'gaussian', 'box', 'boxSpectral', 19 | 'spectral' or 'spectral-circle', 'spectral-square', 'None' 20 | The default is 'gaussian'. 21 | 22 | coarseGrainType : str, optional 23 | The method of coarse graining to use. It can be 'spectral' or 'none'. 24 | The default is 'spectral', which means coarse graining will be done in spectral space. 25 | 26 | Delta : float, optional 27 | The characteristic scale of the filter. If not provided, it will be computed from the provided Ngrid. 28 | 29 | Ngrid : list or tuple, optional 30 | The grid sizes in x and y directions. Must be provided if Delta is not provided. 31 | 32 | spectral : bool, optional 33 | If True, the input data U is considered to be already Fourier-transformed. 34 | The default is False, in which case U is Fourier-transformed within the function. 35 | 36 | Returns 37 | ------- 38 | numpy.ndarray 39 | The filtered and coarse-grained version of the input data U. 40 | """ 41 | 42 | # Fourier transform the input data if not already done 43 | if not spectral: 44 | U_hat = np.fft.rfft2(U) 45 | else: 46 | U_hat = U 47 | 48 | # Get grid size in x and y directions 49 | NX_DNS = np.shape(U_hat)[0] 50 | Lx, Ly = 2 * np.pi, 2 * np.pi # Domain size 51 | 52 | # If Delta is not provided, compute it from Ngrid 53 | if Delta is None: 54 | if Ngrid is None: 55 | raise ValueError("Must provide either Delta or Ngrid") 56 | else: 57 | Delta = 2 * Lx / Ngrid[0] 58 | 59 | # Initialize wavenumbers for the DNS grid 60 | Kx_DNS, Ky_DNS, _, Ksq_DNS, _ = initialize_wavenumbers_rfft2(NX_DNS, NX_DNS, Lx, Ly, INDEXING='ij') 61 | 62 | # Apply filter to the data 63 | if filterType == 'gaussian': 64 | Gk = np.exp(-Ksq_DNS * (Delta ** 2) / 24) 65 | U_f_hat = Gk * U_hat 66 | 67 | elif filterType in ['box', 'boxSpectral']: 68 | Gkx = np.sinc(0.5 * Kx_DNS * Delta / np.pi) # numpy's sinc includes pi factor 69 | Gky = np.sinc(0.5 * Ky_DNS * Delta / np.pi) 70 | Gk = Gkx * Gky 71 | U_f_hat = Gk * U_hat 72 | 73 | elif filterType == 'spectral' or filterType == 'spectral-circle': 74 | kc = Lx / (Delta) 75 | U_f_hat = spectral_filter_circle_same_size(U_hat, kc) 76 | 77 | elif filterType == 'spectral-square': 78 | kc = Lx / (Delta) 79 | U_f_hat = spectral_filter_square_same_size(U_hat, kc) 80 | 81 | elif filterType == None: 82 | U_f_hat = U_hat 83 | 84 | else: 85 | raise ValueError("Invalid filter type : " + filterType) 86 | 87 | # Apply coarse graining 88 | if coarseGrainType == 'spectral': 89 | U_f_c_hat = coarse_spectral_filter_square(U_f_hat, Ngrid[0]) 90 | 91 | elif coarseGrainType == None: 92 | U_f_c_hat = U_f_hat 93 | 94 | else: 95 | raise ValueError("Invalid coarse grain type : " + coarseGrainType) 96 | 97 | # Inverse Fourier transform the result and return the real part 98 | if not spectral: 99 | return np.fft.irfft2(U_f_c_hat, s=(U_f_c_hat.shape[0], U_f_c_hat.shape[0])) 100 | else: 101 | return U_f_c_hat 102 | 103 | 104 | def invertFilter2D(Uf, filterType='gaussian', Delta=None, spectral=False): 105 | """ 106 | Invert filters (only for valid for invertible filters) used in 2D Turbulence. 107 | 108 | Parameters 109 | ---------- 110 | U : numpy.ndarray 111 | The 2D input data to be filtered and coarse grained. 112 | 113 | filterType : str, optional 114 | The type of filter to apply. It can be 'gaussian', 'box' 115 | The default is 'gaussian'. 116 | 117 | 118 | Delta : float, optional 119 | The characteristic scale of the filter. If not provided, it will be computed from the provided Ngrid. 120 | 121 | spectral : bool, optional 122 | If True, the input data U is considered to be already Fourier-transformed. 123 | The default is False, in which case U is Fourier-transformed within the function. 124 | 125 | Returns 126 | ------- 127 | numpy.ndarray 128 | The inverted data 129 | """ 130 | 131 | # Fourier transform the input data if not already done 132 | if not spectral: 133 | Uf_hat = np.fft.rfft2(Uf) 134 | else: 135 | Uf_hat = Uf 136 | 137 | # Get grid size in x and y directions 138 | Ngrid = np.shape(Uf_hat)[0] 139 | Lx, Ly = 2 * np.pi, 2 * np.pi # Domain size 140 | 141 | # If Delta is not provided, compute it from Ngrid 142 | if Delta is None: 143 | Delta = 2 * Lx / Ngrid 144 | 145 | # Initialize wavenumbers for the DNS grid 146 | Kx, Ky, _, Ksq, _ = initialize_wavenumbers_rfft2(Ngrid, Ngrid, Lx, Ly, INDEXING='ij') 147 | 148 | # Apply filter to the data 149 | if filterType == 'gaussian': 150 | Gk = np.exp(-Ksq * (Delta ** 2) / 24) 151 | 152 | elif filterType in ['box', 'boxSpectral']: 153 | Gkx = np.sinc(0.5 * Kx * Delta / np.pi) # numpy's sinc includes pi factor 154 | Gky = np.sinc(0.5 * Ky * Delta / np.pi) 155 | Gkx[Ngrid//2, :] = 1.0 # these values are zero; avoiding division by zero in next steps 156 | Gky[:,Ngrid//2] = 1.0 157 | Gk = Gkx * Gky 158 | 159 | else: 160 | raise ValueError("Invalid filter type : " + filterType) 161 | 162 | U_hat = Uf_hat/Gk 163 | 164 | # Inverse Fourier transform the result and return the real part 165 | if not spectral: 166 | return np.fft.irfft2(U_hat, s=(U_hat.shape[0], U_hat.shape[0])) 167 | else: 168 | return U_hat 169 | 170 | 171 | def spectral_filter_circle_same_size(q_hat, kc): 172 | ''' 173 | A circular sharp spectral filter for 2D flow variables. The function takes a 2D square matrix at high resolution 174 | and Ngrid_coarse for low resolution. The function coasre grains the data in spectral domain and 175 | returns the low resolution data in the frequency domain. 176 | 177 | Parameters: 178 | q (numpy.ndarray): The input 2D square matrix. 179 | kc (int): The cutoff wavenumber. 180 | 181 | Returns: 182 | numpy.ndarray: The filtered data. The data is in the frequency domain. 183 | ''' 184 | NX_DNS = q_hat.shape[0] 185 | Lx, Ly = 2 * np.pi, 2 * np.pi 186 | 187 | _, _, Kabs_DNS, _, _ = initialize_wavenumbers_rfft2(NX_DNS, NX_DNS, Lx, Ly, INDEXING='ij') 188 | 189 | q_filtered_hat = np.where(Kabs_DNS < kc, q_hat, 0) 190 | return q_filtered_hat 191 | 192 | def spectral_filter_square_same_size(q_hat, kc): 193 | ''' 194 | A sharp spectral filter for 2D flow variables. The function takes a 2D square matrix at high resolution 195 | The function applies a square sharp-spectral (cut-off filter) to the input data in the frequency domain. 196 | 197 | Parameters: 198 | q (numpy.ndarray): The input 2D square matrix. 199 | kc (int): The cutoff wavenumber. 200 | 201 | Returns: 202 | numpy.ndarray: The filtered data. The data is in the frequency domain. 203 | ''' 204 | NX_DNS = q_hat.shape[0] 205 | Lx, Ly = 2 * np.pi, 2 * np.pi 206 | 207 | Kx_DNS, Ky_DNS, _, _, _ = initialize_wavenumbers_rfft2(NX_DNS, NX_DNS, Lx, Ly, INDEXING='ij') 208 | 209 | Kx_DNS_abs = np.abs(Kx_DNS) 210 | Ky_DNS_abs = np.abs(Ky_DNS) 211 | 212 | q_filtered_hat = np.where((Kx_DNS_abs < kc) & (Ky_DNS_abs < kc), q_hat, 0) 213 | return q_filtered_hat 214 | 215 | def coarse_spectral_filter_square(a_hat, NCoarse): 216 | """ 217 | Apply a coarse spectral filter to the Fourier-transformed input on a 2D square grid. 218 | 219 | This function filters the Fourier space representation of some input data, effectively removing 220 | high-frequency information above a certain threshold determined by the number of effective 221 | large eddy simulation (LES) points, `NLES`. The filter is applied on a square grid in Fourier space. 222 | 223 | Parameters 224 | ---------- 225 | a_hat : numpy.ndarray 226 | The 2D Fourier-transformed input data, expected to be a square grid. 227 | 228 | NLES : int 229 | The number of effective LES points, determining the cutoff for the spectral filter. 230 | Frequencies beyond half of this value will be cut off. 231 | 232 | Returns 233 | ------- 234 | numpy.ndarray 235 | The filtered Fourier-transformed data. 236 | 237 | Note 238 | ---- 239 | Works for both odd and even number of grid points 240 | """ 241 | 242 | # Check if both original grid and coarse grid are even or odd 243 | if np.mod(a_hat.shape[0], 2) != np.mod(NCoarse, 2): 244 | print(a_hat.shape, NCoarse) 245 | raise ValueError("Both original grid and coarse grid should be even or odd") 246 | 247 | 248 | # Determine the size of the input array 249 | N = a_hat.shape[0] 250 | 251 | # Index of Nyquist wave number - fine grid 252 | kn_fine = N//2 253 | 254 | # Index of Nyquist wave number - coarse grid 255 | kn_coarse =int(NCoarse//2) 256 | 257 | # Compute indices for coarse grid 258 | indvpadx = np.r_[0:kn_coarse+1, kn_fine+1+(kn_fine-kn_coarse):N] 259 | indvpady = np.r_[0:kn_coarse+1] 260 | 261 | # Apply scaling and pad the spectral data with zeros 262 | u_hat_coarse = a_hat[np.ix_(indvpadx, indvpady)]*((NCoarse/N)**2) 263 | 264 | # ################## Making the data conjugate symmetric ################## 265 | 266 | if NCoarse % 2 == 0: 267 | 268 | # # # ** Method#1: Making the array conjugate symmetric 269 | u_hat_coarse[kn_coarse,0] = np.real(u_hat_coarse[kn_coarse,0]) 270 | u_hat_coarse[0,kn_coarse] = np.real(u_hat_coarse[0,kn_coarse]) 271 | u_hat_coarse[1:,kn_coarse] = 0.5*(u_hat_coarse[1:,kn_coarse] + np.conj(np.flip(u_hat_coarse[1:,kn_coarse]))) 272 | 273 | # # # ** Method #2: Making the Nyquist wavenumber 0 274 | # u_hat_coarse[kn_coarse,:] = 0 275 | # u_hat_coarse[:,kn_coarse] = 0 276 | pass 277 | 278 | else: # Odd grid size 279 | 280 | # ** Method#1: Making the array conjugate symmetric 281 | u_hat_coarse[0,kn_coarse] = np.real(u_hat_coarse[0,kn_coarse]) 282 | 283 | # # ** Method #2: Making the Nyquist wavenumber 0 284 | # u_hat_coarse[kn_coarse,:] = 0 285 | # u_hat_coarse[kn_coarse+1,:] = 0 286 | # u_hat_coarse[:,kn_coarse] = 0 287 | 288 | 289 | # ** Method#3: take irfft and back: works for both odd and even grid sizes: Does same thing as Method#1 290 | u = np.fft.irfft2(u_hat_coarse, s=[NCoarse,NCoarse]) 291 | u_hat_coarse = np.fft.rfft2(u) 292 | 293 | return u_hat_coarse 294 | 295 | # # Code deals with ifft's and fft's 296 | # def coarse_spectral_filter_square_2DFHIT(a_hat, NCoarse): 297 | # """ 298 | # Apply a coarse spectral filter to the Fourier-transformed input on a 2D square grid. 299 | 300 | # This function filters the Fourier space representation of some input data, effectively removing 301 | # high-frequency information above a certain threshold determined by the number of effective 302 | # large eddy simulation (LES) points, `NLES`. The filter is applied on a square grid in Fourier space. 303 | 304 | # Parameters 305 | # ---------- 306 | # a_hat : numpy.ndarray 307 | # The 2D Fourier-transformed input data, expected to be a square grid. 308 | 309 | # NLES : int 310 | # The number of effective LES points, determining the cutoff for the spectral filter. 311 | # Frequencies beyond half of this value will be cut off. 312 | 313 | # Returns 314 | # ------- 315 | # numpy.ndarray 316 | # The filtered Fourier-transformed data. 317 | # """ 318 | 319 | # # Determine the size of the input array 320 | # N = a_hat.shape[0] 321 | 322 | # # Compute the cutoff point in Fourier space 323 | # dkcut= int(NCoarse/2) 324 | 325 | # # Define the start and end indices for the slice in Fourier space to keep 326 | # ids = int(N/2)-dkcut 327 | # ide = int(N/2)+dkcut 328 | 329 | # # Shift the zero-frequency component to the center, then normalize the Fourier-transformed data 330 | # a_hat_shift = np.fft.fftshift(a_hat)/(N**2) 331 | 332 | # # Apply the spectral filter by slicing the 2D array 333 | # wfiltered_hat_shift = a_hat_shift[ids:ide,ids:ide] 334 | 335 | # # Shift the zero-frequency component back to the original place and un-normalize the data 336 | # wfiltered_hat = np.fft.ifftshift(wfiltered_hat_shift)*(NCoarse**2) 337 | 338 | # # Ensure Nyquist wavenumber is its own complex conjugate 339 | # wfiltered_hat_sym = conjugate_symmetrize_coarse(wfiltered_hat) 340 | 341 | # # Return the filtered data 342 | # return wfiltered_hat_sym 343 | 344 | # def conjugate_symmetrize_coarse(a_hat): 345 | # # Ensures Nyquist wavenumbers are complex conjugates 346 | # # This function is to be used with square_spectral_coarse_grain 347 | 348 | # NCoarse = int(a_hat.shape[0]) 349 | # a_hat_sym = a_hat.copy() 350 | 351 | # # # 0th Wavenumber is conjugate symmetric 352 | # a_hat_sym[0,NCoarse//2+1:] = np.flip(a_hat[0,1:NCoarse//2]).conj() 353 | # a_hat_sym[NCoarse//2+1:,0] = np.flip(a_hat[1:NCoarse//2,0]).conj() 354 | 355 | # # # Nyquist frequency is conjugate symmetric 356 | # a_hat_sym[NCoarse//2,NCoarse//2+1:] = np.flip(a_hat[NCoarse//2,1:NCoarse//2]).conj() 357 | # a_hat_sym[NCoarse//2+1:,NCoarse//2] = np.flip(a_hat[1:NCoarse//2,NCoarse//2]).conj() 358 | 359 | # # # 0th wavenumber 0th wavenumbers has zero imaginary part 360 | # a_hat_sym[0,0] = a_hat[0,0].real # (Kx=0, Ky=0) 361 | 362 | # # # Nyquist wavenumber of 0th wavenumber has zero imaginary part 363 | # a_hat_sym[0,NCoarse//2] = a_hat[0,NCoarse//2].real # (Kx=0, Ky=Ny/2) 364 | # a_hat_sym[NCoarse//2,0] = a_hat[NCoarse//2,0].real # (Kx=Nx/2, Ky=0) 365 | 366 | # # # Nyquist frequency of Nyquist frequency has zero imaginary part 367 | # a_hat_sym[NCoarse//2,NCoarse//2] = a_hat[NCoarse//2,NCoarse//2].real # (Kx=Nx/2, Ky=Ny/2) 368 | 369 | # ########################## 370 | # # #Ensure Nyquist wavenumber is its own complex conjugate 371 | # # a_hat_sym[NCoarse//2,1:] = np.real((a_hat[NCoarse//2,1:] + np.conj(np.flip(a_hat[NCoarse//2,1:])))/2) 372 | # # a_hat_sym[1:,NCoarse//2] = np.real((a_hat[1:,NCoarse//2] + np.conj(np.flip(a_hat[1:,NCoarse//2])))/2) 373 | 374 | # # ##### Remove data at the nyquist frequency to make it conjugate frequency #####) 375 | # a_hat_sym[NCoarse//2, :] = 0 376 | # a_hat_sym[:,NCoarse//2] = 0 377 | 378 | # return a_hat_sym 379 | 380 | # # JAX compatible 381 | # def coarse_spectral_filter_square_2DFHIT_jit(a_hat, NCoarse): 382 | # """ 383 | # Apply a coarse spectral filter to the Fourier-transformed input on a 2D square grid. 384 | 385 | # This function filters the Fourier space representation of some input data, effectively removing 386 | # high-frequency information above a certain threshold determined by the number of effective 387 | # large eddy simulation (LES) points, `NLES`. The filter is applied on a square grid in Fourier space. 388 | 389 | # Parameters 390 | # ---------- 391 | # a_hat : numpy.ndarray 392 | # The 2D Fourier-transformed input data, expected to be a square grid. 393 | 394 | # NLES : int 395 | # The number of effective LES points, determining the cutoff for the spectral filter. 396 | # Frequencies beyond half of this value will be cut off. 397 | 398 | # Returns 399 | # ------- 400 | # numpy.ndarray 401 | # The filtered Fourier-transformed data. 402 | # """ 403 | 404 | # # Determine the size of the input array 405 | # N = a_hat.shape[0] 406 | 407 | # # Compute the cutoff point in Fourier space 408 | # dkcut= NCoarse//2 409 | 410 | # # Shift the zero-frequency component to the center, then normalize the Fourier-transformed data 411 | # a_hat_shift = jnp.fft.fftshift(a_hat)/(N**2) 412 | 413 | # # Define the start and end indices for the slice in Fourier space to keep 414 | # ids = N//2-dkcut 415 | # ide = N//2+dkcut 416 | 417 | # # Apply the spectral filter by slicing the 2D array 418 | # wfiltered_hat_shift = a_hat_shift[ids:ide,ids:ide] 419 | 420 | # # Shift the zero-frequency component back to the original place and un-normalize the data 421 | # wfiltered_hat = jnp.fft.ifftshift(wfiltered_hat_shift)*(NCoarse**2) 422 | 423 | # # Ensure Nyquist wavenumber is its own complex conjugate 424 | # wfiltered_hat_sym = conjugate_symmetrize_coarse_jit(wfiltered_hat) 425 | 426 | # # Return the filtered data 427 | # return wfiltered_hat_sym 428 | 429 | # @jit 430 | # def conjugate_symmetrize_coarse_jit(a_hat_coarse): 431 | # # Ensures Nyquist wavenumbers are complex conjugates 432 | # # This function is to be used with square_spectral_coarse_grain 433 | 434 | # NCoarse = int(a_hat_coarse.shape[0]) 435 | # a_hat_coarse_sym = a_hat_coarse.copy() 436 | 437 | # # #Ensure Nyquist wavenumber is its own complex conjugate 438 | # # a_hat_coarse_sym = a_hat_coarse_sym.at[NCoarse//2, 0].set(jnp.real(a_hat_coarse_sym[NCoarse//2, 0])) 439 | # # a_hat_coarse_sym = a_hat_coarse_sym.at[0, NCoarse//2].set(jnp.real(a_hat_coarse_sym[0, NCoarse//2])) 440 | # # a_hat_coarse_sym = a_hat_coarse_sym.at[NCoarse//2,1:].set((a_hat_coarse_sym[NCoarse//2,1:] + jnp.conj(jnp.flip(a_hat_coarse_sym[NCoarse//2,1:])))/2) 441 | # # a_hat_coarse_sym = a_hat_coarse_sym.at[1:,NCoarse//2].set((a_hat_coarse_sym[1:,NCoarse//2] + jnp.conj(jnp.flip(a_hat_coarse_sym[1:,NCoarse//2])))/2) 442 | 443 | # ##### Alternatively remove the data (make it zero) at the Nyquist frequency to make it conjugate symmetric ##### 444 | # a_hat_coarse_sym = a_hat_coarse_sym.at[NCoarse//2, :].set(0) 445 | # a_hat_coarse_sym = a_hat_coarse_sym.at[:, NCoarse//2].set(0) 446 | 447 | # return a_hat_coarse_sym 448 | 449 | # @jit 450 | def filter2D_gaussian_spectral_jit(U_hat, Ksq, Delta): 451 | 452 | Gk = jnp.exp(-Ksq * (Delta ** 2) / 24) 453 | U_f_hat = Gk * U_hat 454 | 455 | return U_f_hat 456 | 457 | @jit 458 | def filter2D_box_spectral_jit(U_hat, Kx, Ky, Delta): 459 | 460 | Gkx = jnp.sinc(0.5 * Kx * Delta / jnp.pi) 461 | Gky = jnp.sinc(0.5 * Ky * Delta / jnp.pi) 462 | 463 | Gkx = Gkx.at[0, :].set(1.0) 464 | Gky = Gky.at[:, 0].set(1.0) 465 | 466 | Gk = Gkx * Gky 467 | 468 | U_f_hat = Gk * U_hat 469 | 470 | return U_f_hat 471 | 472 | # @jit 473 | def invertFilter2D_gaussian_spectral_jit(Uf_hat, Ksq, Delta): 474 | Gk = jnp.exp(-Ksq * (Delta ** 2) / 24) 475 | U_hat = Uf_hat/Gk 476 | 477 | return U_hat 478 | 479 | @jit 480 | def invertFilter2D_box_spectral_jit(Uf_hat, Kx, Ky, Delta): 481 | Gkx = jnp.sinc(0.5 * Kx * Delta / jnp.pi) # jnp.sinc includes pi factor 482 | Gky = jnp.sinc(0.5 * Ky * Delta / jnp.pi) 483 | 484 | Gkx = Gkx.at[0, :].set(1.0) 485 | Gky = Gky.at[:, 0].set(1.0) 486 | 487 | Ngrid = Uf_hat.shape[0] 488 | Gkx = Gkx.at[Ngrid//2, :].set(1.0) # avoiding division by zero in next steps 489 | Gky = Gky.at[:, Ngrid//2].set(1.0) 490 | 491 | Gk = Gkx * Gky 492 | 493 | U_hat = Uf_hat / Gk 494 | 495 | return U_hat 496 | 497 | @jit 498 | def spectral_filter_square_same_size_jit(q_hat, kc): 499 | ''' 500 | A sharp spectral filter for 2D flow variables. The function takes a 2D square matrix at high resolution 501 | The function applies a square sharp-spectral (cut-off filter) to the input data in the frequency domain. 502 | 503 | Parameters: 504 | q (numpy.ndarray): The input 2D square matrix. 505 | kc (int): The cutoff wavenumber. 506 | 507 | Returns: 508 | numpy.ndarray: The filtered data. The data is in the frequency domain. 509 | ''' 510 | NX_DNS = q_hat.shape[0] 511 | Lx, Ly = 2 * jnp.pi, 2 * jnp.pi 512 | 513 | # Create an array of the discrete Fourier Transform sample frequencies in x,y-direction 514 | kx = 2 * jnp.pi * jnp.fft.fftfreq(NX_DNS, d=Lx/NX_DNS) 515 | ky = 2 * jnp.pi * jnp.fft.fftfreq(NX_DNS, d=Ly/NX_DNS) 516 | 517 | # Return coordinate grids (2D arrays) for the x and y wavenumbers 518 | (Kx_DNS, Ky_DNS) = jnp.meshgrid(kx, ky, indexing='ij') 519 | 520 | Kx_DNS_abs = jnp.abs(fft2_to_rfft2(Kx_DNS)) 521 | Ky_DNS_abs = jnp.abs(fft2_to_rfft2(Ky_DNS)) 522 | 523 | q_filtered_hat = jnp.where((Kx_DNS_abs < kc) & (Ky_DNS_abs < kc), q_hat, 0) 524 | return q_filtered_hat 525 | 526 | def coarse_spectral_filter_square_jit(a_hat, NCoarse): 527 | """ 528 | Apply a coarse spectral filter to the Fourier-transformed input on a 2D square grid. 529 | 530 | This function filters the Fourier space representation of input data, removing 531 | high-frequency information above a threshold determined by the number of effective 532 | large eddy simulation (LES) points, `NCoarse`. 533 | 534 | Args: 535 | a_hat: JAX array. The 2D Fourier-transformed input data (square grid). 536 | NLES: int. The number of effective LES points, defining the spectral filter cutoff. 537 | 538 | Returns: 539 | JAX array. The filtered Fourier-transformed data. 540 | 541 | Note: 542 | Works for both odd and even number of grid points 543 | """ 544 | 545 | # Check if shapes are compatible 546 | # if jnp.mod(a_hat.shape[0], 2) != jnp.mod(NCoarse, 2): 547 | # raise ValueError("Both original grid and coarse grid should be even or odd") 548 | 549 | N = a_hat.shape[0] 550 | kn_fine = N // 2 551 | kn_coarse = NCoarse // 2 552 | 553 | # Apply scaling coarse graining 554 | u_hat_coarse = jnp.zeros((NCoarse, NCoarse//2+1), dtype=complex) 555 | u_hat_coarse = u_hat_coarse.at[:kn_coarse+1, :].set(a_hat[:kn_coarse+1, :kn_coarse+1] ) 556 | u_hat_coarse = u_hat_coarse.at[kn_coarse+1:, :].set(a_hat[kn_fine+1+(kn_fine-kn_coarse):, :kn_coarse+1] ) 557 | u_hat_coarse = u_hat_coarse * ((NCoarse / N) ** 2) 558 | 559 | # Enforce conjugate symmetry 560 | if NCoarse % 2 == 0: 561 | # # Method #1: Making the array conjugate symmetric 562 | # u_hat_coarse = u_hat_coarse.at[kn_coarse, 0].set(jnp.real(u_hat_coarse[kn_coarse, 0])) 563 | # u_hat_coarse = u_hat_coarse.at[0, kn_coarse].set(jnp.real(u_hat_coarse[0, kn_coarse])) 564 | # u_hat_coarse = u_hat_coarse.at[1:, kn_coarse].set( 565 | # 0.5 * (u_hat_coarse[1:, kn_coarse] + jnp.conj(jnp.flip(u_hat_coarse[1:, kn_coarse]))) 566 | # ) 567 | 568 | # Method #2: Making the Nyquist wavenumber 0 569 | u_hat_coarse = u_hat_coarse.at[kn_coarse,:].set(0) 570 | u_hat_coarse = u_hat_coarse.at[:,kn_coarse].set(0) 571 | 572 | else: # Odd grid size 573 | # Method #1: Making the array conjugate symmetric 574 | # u_hat_coarse = u_hat_coarse.at[0,kn_coarse].set(jnp.real(u_hat_coarse[0,kn_coarse])) 575 | 576 | # # Method #2: Making the Nyquist wavenumber 0 577 | u_hat_coarse = u_hat_coarse.at[kn_coarse, :].set(0.0) 578 | u_hat_coarse = u_hat_coarse.at[kn_coarse + 1, :].set(0.0) 579 | u_hat_coarse = u_hat_coarse.at[:, kn_coarse].set(0.0) 580 | 581 | # # Method #3: take irfft and back: works for both odd and even grid sizes 582 | # u = jnp.fft.irfft2(u_hat_coarse, s=[NCoarse,NCoarse]) 583 | # u_hat_coarse = jnp.fft.rfft2(u) 584 | 585 | return u_hat_coarse 586 | 587 | -------------------------------------------------------------------------------- /py2d/initialize.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from py2d.util import fft2_to_rfft2 3 | 4 | def initialize_perturbation(NX, Kx, Ky): 5 | # -------------- Initialization using pertubration -------------- 6 | kp = 10.0 7 | A = 4*np.power(kp,(-5))/(3*np.pi) 8 | absK = np.sqrt(Kx*Kx+Ky*Ky) 9 | Ek = A*np.power(absK,4)*np.exp(-np.power(absK/kp,2)) 10 | coef1 = np.random.uniform(0,1,[NX//2+1,NX//2+1])*np.pi*2 11 | coef2 = np.random.uniform(0,1,[NX//2+1,NX//2+1])*np.pi*2 12 | 13 | perturb = np.zeros([NX,NX]) 14 | perturb[0:NX//2+1, 0:NX//2+1] = coef1[0:NX//2+1, 0:NX//2+1]+coef2[0:NX//2+1, 0:NX//2+1] 15 | perturb[NX//2+1:, 0:NX//2+1] = coef2[NX//2-1:0:-1, 0:NX//2+1] - coef1[NX//2-1:0:-1, 0:NX//2+1] 16 | perturb[0:NX//2+1, NX//2+1:] = coef1[0:NX//2+1, NX//2-1:0:-1] - coef2[0:NX//2+1, NX//2-1:0:-1] 17 | perturb[NX//2+1:, NX//2+1:] = -(coef1[NX//2-1:0:-1, NX//2-1:0:-1] + coef2[NX//2-1:0:-1, NX//2-1:0:-1]) 18 | perturb = np.exp(1j*perturb) 19 | 20 | w1_hat = np.sqrt(absK/np.pi*Ek)*perturb*np.power(NX,2) 21 | 22 | psi_hat = -w1_hat*invKsq 23 | psiPrevious_hat = psi_hat.astype(np.complex128) 24 | psiCurrent_hat = psi_hat.astype(np.complex128) 25 | 26 | return w1_hat, psi_hat, psiPrevious_hat, psiCurrent_hat 27 | 28 | 29 | def initialize_wavenumbers_fft2(nx, ny, Lx, Ly, INDEXING='ij'): 30 | ''' 31 | Initialize the wavenumbers for 2D Forced Homogeneous Isotropic Turbulence (2D-FHIT). 32 | 33 | Parameters: 34 | ----------- 35 | nx : int 36 | Number of grid points in the x-direction. 37 | ny : int 38 | Number of grid points in the y-direction. 39 | Lx : float 40 | Length of the domain in the x-direction. 41 | Ly : float 42 | Length of the domain in the y-direction. 43 | 44 | Returns: 45 | -------- 46 | Kx : numpy.ndarray 47 | 2D array of wavenumbers in the x-direction. 48 | Ky : numpy.ndarray 49 | 2D array of wavenumbers in the y-direction. 50 | Kabs : numpy.ndarray 51 | 2D array of the absolute values of the wavenumbers. 52 | Ksq : numpy.ndarray 53 | 2D array of the square of the wavenumber magnitudes. 54 | invKsq : numpy.ndarray 55 | 2D array of the inverse of the square of the wavenumber magnitudes. 56 | 57 | 58 | Notes: 59 | ------ 60 | inKsq[0,0] = 0 to avoid numerical errors, since invKsq[0.0] = 1/0 = inf 61 | ''' 62 | 63 | # Create an array of the discrete Fourier Transform sample frequencies in x-direction 64 | kx = 2 * np.pi * np.fft.fftfreq(nx, d=Lx/nx) 65 | 66 | # Create an array of the discrete Fourier Transform sample frequencies in y-direction 67 | ky = 2 * np.pi * np.fft.fftfreq(ny, d=Ly/ny) 68 | 69 | # Return coordinate grids (2D arrays) for the x and y wavenumbers 70 | (Kx, Ky) = np.meshgrid(kx, ky, indexing=INDEXING) 71 | 72 | # Compute the squared magnitudes of the 2D wavenumbers (Kx and Ky) 73 | Ksq = Kx ** 2 + Ky ** 2 74 | 75 | # Compute the absolute value of the wavenumbers 76 | Kabs = np.sqrt(Ksq) 77 | 78 | # To avoid division by zero, set the zero wavenumber to a large value 79 | Ksq[0,0] = 1e16 80 | 81 | # Compute the inverse of the squared wavenumbers 82 | invKsq = 1.0 / Ksq 83 | # Set the inverse of the zero wavenumber to zero 84 | invKsq[0,0] = 0.0 85 | 86 | # Set the zero wavenumber back to zero 87 | Ksq[0,0] = 0.0 88 | 89 | # Return the wavenumbers in the x and y directions, their absolute values, 90 | # their squared magnitudes and inverse of the squared magnitudes 91 | return Kx, Ky, Kabs, Ksq, invKsq 92 | 93 | def initialize_wavenumbers_rfft2(nx, ny, Lx, Ly, INDEXING='ij'): 94 | 95 | Kx, Ky, Kabs, Ksq, invKsq = initialize_wavenumbers_fft2(nx, ny, Lx, Ly, INDEXING=INDEXING) 96 | 97 | return fft2_to_rfft2(Kx), fft2_to_rfft2(Ky), fft2_to_rfft2(Kabs), fft2_to_rfft2(Ksq), fft2_to_rfft2(invKsq) 98 | 99 | 100 | def gridgen(Lx, Ly, Nx, Ny, INDEXING='ij'): 101 | ''' 102 | Generate a 2D grid. 103 | 104 | Parameters: 105 | ----------- 106 | Lx : float 107 | Length of the domain in the x-direction. 108 | NX : int 109 | Number of grid points in the x and y-directions. 110 | INDEXING : str, optional 111 | Convention to use for indexing. Default is 'ij' (matrix indexing). 112 | 113 | Returns: 114 | -------- 115 | Lx : float 116 | Length of the domain in the x-direction. 117 | Lx : float 118 | Length of the domain in the y-direction (same as x-direction as grid is square). 119 | X : numpy.ndarray 120 | 2D array of x-coordinates. 121 | Y : numpy.ndarray 122 | 2D array of y-coordinates. 123 | dx : float 124 | Size of grid spacing in the x-direction. 125 | dx : float 126 | Size of grid spacing in the y-direction (same as x-direction as grid is square). 127 | ''' 128 | 129 | # Calculate the size of the grid spacing 130 | dx = Lx / Nx 131 | dy = Ly / Ny 132 | 133 | # Create an array of x-coordinates, ranging from 0 to (Lx - dx) 134 | x = np.linspace(0, Lx - dx, num=Nx) 135 | y = np.linspace(0, Lx - dx, num=Ny) 136 | 137 | # Create 2D arrays of the x and y-coordinates using a meshgrid. 138 | X, Y = np.meshgrid(x, y, indexing=INDEXING) 139 | 140 | # Return the lengths of the domain, the x and y-coordinates, and the size of the grid spacing. 141 | return Lx, Ly, X, Y, dx, dy 142 | -------------------------------------------------------------------------------- /py2d/readme.md: -------------------------------------------------------------------------------- 1 | # SGS models 2 | 3 | ## Eddy viscosity models 4 | 5 | $$ \nabla . \tau = \nabla.(\nu_e \nabla \omega )$$ 6 | 7 | 8 | ### Eddy viscosity, Smagorinsky 9 | ### Eddy viscosity, Leith 10 | 11 | ### Eddy viscosity, Dnamic Smagorinsky 12 | ### Eddy viscosity, Dynamic Leith 13 | 14 | ## Local models 15 | 16 | 17 | $$\Pi = ∇.(ν_e ∇ω )$$ 18 | 19 | versus 20 | 21 | $$\Pi= ν_e ∇.( ∇ω ) == ν_e (∇^2ω )$$ 22 | 23 | 24 | ![pi](../assets/7940296/bd635cb3-ca3e-497c-9eb8-d032079cdb37) 25 | 26 | Options 27 | 28 | 1. Domain Averaged $C_s$ 29 | 30 | $$ 31 | \nu_e(x,t) = ( C_s(x,t) \Delta )^2 [ 2 |\bar{S}(x,t) |^2 ]^{(1/2)} 32 | $$ 33 | 34 | 2. Local $C_s$ 35 | 36 | $$ 37 | \nu_e(x,t) = ( C_s(x,t) \Delta )^2 [ 2 |\bar{S}(x,t) |^2 ]^{(1/2)} 38 | $$ 39 | 40 | 3. Domain average 41 | 42 | 43 | From Pierre Sagaut [1] 44 | 45 | 46 | ![sagaut](../assets/7940296/a488d7ed-c320-46b7-a2c5-c78bd37bcdde) 47 | 48 | 49 | [1] 50 | -------------------------------------------------------------------------------- /py2d/spectra.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from py2d.initialize import initialize_wavenumbers_fft2 4 | from py2d.derivative import derivative 5 | 6 | def spectrum_angled_average(A, kmax = 'grid', spectral = False): 7 | ''' 8 | Compute the radially/angle-averaged spectrum of a 2D square matrix. 9 | 10 | Parameters 11 | ---------- 12 | A : numpy.ndarray 13 | The input 2D square matrix. If `spectral` is False, `A` is in the physical domain; otherwise, it is in the spectral domain. 14 | kmax : str, optional 'grid' or 'diagonal' Default is 'grid' 15 | The maximum wavenumber to be considered in the average. Default is Ngrid/2 which will consider all wavenumber of radius less than Ngrid/2. 16 | If confused used default. 17 | Alternatively, it will consider all wavenumber upto the radius less than the diagonal of the square grid can with i.e. kmax = np.sqrt(2)*Ngrid/2 18 | spectral : bool, optional 19 | Whether `A` is in the spectral domain. Default is False. 20 | 21 | Returns 22 | ------- 23 | spectrum : numpy.ndarray 24 | The radially/angle-averaged spectrum of `A`. 25 | wavenumbers : numpy.ndarray 26 | The corresponding wavenumbers. 27 | 28 | Raises 29 | ------ 30 | ValueError 31 | If `A` is not a 2D square matrix or `spectral` is not a boolean or `A` contains non-numeric values. 32 | 33 | Notes 34 | ----- 35 | To calculate angled-averaged spectrum of absolute value of a complex-valued matrix, input np.sqrt(np.conj(A_hat)*A_hat) as A. 36 | np.abs(A_hat) is not the same as np.sqrt(np.conj(A_hat)*A_hat) for complex-valued A_hat. 37 | np.abs(A_hat) calculates the magnitude (sqrt(a^2 + b^2)) of the complex-valued A_hat, 38 | while np.sqrt(np.conj(A_hat)*A_hat) calculates the absolute value (sqrt(a^2 + (ib)^2)) of A_hat where A_hat = a + ib. 39 | 40 | 41 | ''' 42 | 43 | # Check if input 'A' is a 2D square matrix 44 | if not isinstance(A, np.ndarray) or A.ndim != 2 : 45 | raise ValueError('Input is not a 2D matrix. Please input a 2D matrix') 46 | # Check if 'spectral' is a boolean value 47 | if not isinstance(spectral, bool): 48 | raise ValueError('Invalid input for spectral. It should be a boolean value') 49 | # Check if input 'A' contains non-numeric values 50 | if not np.issubdtype(A.dtype, np.number): 51 | raise ValueError('Input contains non-numeric values') 52 | 53 | # Calculate the number of grid points in one dimension 54 | nx,ny = A.shape 55 | # Set the size of the domain 56 | Lx, Ly = 2 * np.pi, 2 * np.pi 57 | 58 | _, _, Kabs, _, _ = initialize_wavenumbers_fft2(nx, ny, Lx, Ly, INDEXING='ij') 59 | 60 | # Compute the 2D FFT of 'A' if it is in the physical domain 61 | if not spectral: 62 | A_hat = np.fft.fft2(A) 63 | else: # If 'A' is already in the spectral domain, copy it to 'spectral_A' 64 | A_hat = A 65 | # Normalize the spectrum by the total number of grid points 66 | A_hat = A_hat / nx ** 2 67 | 68 | # Calculate the maximum wavenumber to be considered in the average 69 | if kmax == 'grid': 70 | kxMax = round(nx / 2) 71 | elif kmax == 'diagonal': 72 | kxMax = round(np.sqrt(2) * nx / 2) 73 | 74 | # Initialize the output array with zeros 75 | A_angled_average_spectra = np.zeros(kxMax + 1) 76 | # The zeroth wavenumber is just the first element of the spectrum 77 | A_angled_average_spectra[0] = A_hat[0,0].real 78 | 79 | # Compute the angle-averaged spectrum for wavenumbers 1 to kxMax 80 | for k in range(1, kxMax + 1): 81 | tempInd = (Kabs > (k - 0.5)) & (Kabs <= (k + 0.5)) 82 | A_angled_average_spectra[k] = np.sum(A_hat[tempInd].real) 83 | 84 | # Generate the array of wavenumbers corresponding to the computed spectrum 85 | wavenumbers = np.linspace(0, kxMax, len(A_angled_average_spectra)) 86 | 87 | return A_angled_average_spectra, wavenumbers 88 | 89 | def spectrum_zonal_average(A, spectral=False): 90 | ''' 91 | Compute the zonally-averaged spectrum of a 2D square matrix. 92 | 93 | Parameters 94 | ---------- 95 | A : numpy.ndarray 96 | The input 2D square matrix. If `spectral` is False, `A` is in the physical domain; otherwise, it is in the spectral domain. 97 | spectral : bool, optional 98 | 99 | Returns 100 | ------- 101 | spectrum : numpy.ndarray 102 | The zonally-averaged spectrum of `A`. 103 | wavenumbers : numpy.ndarray 104 | The corresponding wavenumbers. 105 | ''' 106 | 107 | # Check if input 'A' is a 2D square matrix 108 | if not isinstance(A, np.ndarray) or A.ndim != 2 : 109 | raise ValueError('Input is not a 2D matrix. Please input a 2D matrix') 110 | # Check if 'spectral' is a boolean value 111 | if not isinstance(spectral, bool): 112 | raise ValueError('Invalid input for spectral. It should be a boolean value') 113 | 114 | # Calculate the number of grid points in one dimension 115 | nx,ny = A.shape 116 | 117 | # Compute the 1D FFT of 'A' if it is in the physical domain 118 | if not spectral: 119 | A_hat = np.fft.rfft(A, axis=0) 120 | else: # If 'A' is already in the spectral domain, copy it to 'spectral_A' 121 | A_hat = A 122 | 123 | # Normalize the spectrum by the total number of grid points 124 | A_hat = A_hat / nx 125 | 126 | # Compute the zonally-averaged spectrum 127 | A_zonal_average_spectra = np.mean(np.abs(A_hat)*2, axis=1) 128 | 129 | # Generate the array of wavenumbers corresponding to the computed spectrum 130 | wavenumbers = np.linspace(0, nx//2, nx//2+1) 131 | 132 | return A_zonal_average_spectra, wavenumbers 133 | 134 | 135 | def TKE_angled_average(Psi, Omega, spectral=False): 136 | ''' 137 | Calculate the energy spectrum and its angled average. 138 | 139 | Parameters 140 | ---------- 141 | Psi : numpy.ndarray 142 | The input 2D square matrix of stream function. If `spectral` is False, `Psi` is in the physical domain; otherwise, it is in the spectral domain. 143 | Omega : numpy.ndarray 144 | The input 2D square matrix of vorticity. If `spectral` is False, `Omega` is in the physical domain; otherwise, it is in the spectral domain. 145 | spectral : bool, optional 146 | Whether `Psi` and `Omega` are in the spectral domain. Default is False. 147 | 148 | Returns 149 | ------- 150 | TKE_angled_average : numpy.ndarray 151 | The radially/angle-averaged energy spectrum of `E_hat`. 152 | kkx : numpy.ndarray 153 | The corresponding wavenumbers. 154 | ''' 155 | 156 | # Get the size of the input square matrix. 157 | NX = Psi.shape[0] 158 | 159 | # If the data is in physical space, convert it to spectral space using 2D FFT. 160 | if not spectral: 161 | Psi_hat = np.fft.fft2(Psi) 162 | Omega_hat = np.fft.fft2(Omega) 163 | else: # If the data is already in the spectral space, we don't need to transform it. 164 | Psi_hat = Psi 165 | Omega_hat = Omega 166 | 167 | # Compute the kinetic energy in the spectral space. 168 | E_hat = 0.5 * (np.conj(Psi_hat) * Omega_hat) / NX ** 2 169 | 170 | # Perform an angle-averaged spectrum calculation on the computed kinetic energy. 171 | TKE_angled_average, kkx = spectrum_angled_average(E_hat, spectral=True) 172 | 173 | # Return the angle-averaged kinetic energy spectrum and the corresponding wavenumbers. 174 | return TKE_angled_average, kkx 175 | 176 | 177 | def enstrophy_angled_average(Omega, spectral=False): 178 | ''' 179 | Calculate the enstrophy spectrum and its angled average. 180 | 181 | Parameters 182 | ---------- 183 | Omega : numpy.ndarray 184 | The input 2D square matrix of vorticity. If `spectral` is False, `Omega` is in the physical domain; otherwise, it is in the spectral domain. 185 | spectral : bool, optional 186 | Whether `Omega` is in the spectral domain. Default is False. 187 | 188 | Returns 189 | ------- 190 | Z_angled_average_spectra : numpy.ndarray 191 | The radially/angle-averaged enstrophy spectrum of `Z_hat`. 192 | kkx : numpy.ndarray 193 | The corresponding wavenumbers. 194 | ''' 195 | 196 | # Get the size of the input square matrix. 197 | NX = Omega.shape[0] 198 | 199 | # If the data is in physical space, convert it to spectral space using 2D FFT. 200 | if not spectral: 201 | Omega_hat = np.fft.fft2(Omega) 202 | else: # If the data is already in the spectral space, we don't need to transform it. 203 | Omega_hat = Omega 204 | 205 | # Compute the enstrophy in the spectral space. 206 | Z_hat = 0.5 * (np.conj(Omega_hat) * Omega_hat) / NX ** 2 207 | 208 | # Perform an angle-averaged spectrum calculation on the computed enstrophy. 209 | Z_angled_average_spectra, kkx = spectrum_angled_average(Z_hat, spectral=True) 210 | 211 | # Return the angle-averaged enstrophy spectrum and the corresponding wavenumbers. 212 | return Z_angled_average_spectra, kkx 213 | 214 | def energyTransfer_spectra(Kx, Ky, U=None, V=None, Tau11=None, Tau12=None, Tau22=None, Psi=None, PiOmega=None, method='Tau', spectral=False): 215 | ''' 216 | Compute the energy transfer spectra using 2D Forced Homogeneous Isotropic Turbulence (2D-FHIT) 217 | PTau > 0: dissipation 218 | PTau < 0: backscatter 219 | 220 | Parameters 221 | ---------- 222 | Kx, Ky : numpy.ndarray 223 | The 2D wavenumber arrays 224 | U : numpy.ndarray, optional 225 | The U velocity field (2D array) 226 | V : numpy.ndarray, optional 227 | The V velocity field (2D array) 228 | Tau11 : numpy.ndarray, optional 229 | The Tau11 field (2D array) 230 | Tau12 : numpy.ndarray, optional 231 | The Tau12 field (2D array) 232 | Tau22 : numpy.ndarray, optional 233 | The Tau22 field (2D array) 234 | Psi : numpy.ndarray, optional 235 | The Psi field (2D array) 236 | PiOmega : numpy.ndarray, optional 237 | The PiOmega field (2D array) 238 | method : str, optional 239 | The method to calculate Ptau_hat. 'Tau' (default) or 'PiOmega' 240 | spectral : bool, optional 241 | Determines whether the input is in the 'physical' space (default=False) or 'spectral' space 242 | 243 | Returns 244 | ------- 245 | spectra : numpy.ndarray 246 | The computed energy transfer spectra 247 | wavenumber : numpy.ndarray 248 | The corresponding wavenumbers 249 | 250 | This function computes the energy transfer spectra of the given 2D velocity fields 251 | using the 2D Forced Homogeneous Isotropic turbulence. The inputs can be in the physical or spectral space. 252 | ''' 253 | 254 | # If the method is 'Tau', calculate the energy transfer using Tau11, Tau12, Tau22, U, and V. 255 | if method == 'Tau': 256 | # Check if the necessary parameters are provided. 257 | if U is None or V is None or Tau11 is None or Tau12 is None or Tau22 is None: 258 | raise ValueError("U, V, Tau11, Tau12, Tau22 must be provided to calculate energy transfer using 'Tau' method") 259 | 260 | # Calculate the derivatives of U with respect to x and y, and the derivative of V with respect to x. 261 | Ux = derivative(U, [1, 0], Kx=Kx, Ky=Ky, spectral=spectral) 262 | Uy = derivative(U, [0, 1], Kx=Kx,Ky=Ky, spectral=spectral) 263 | Vx = derivative(V, [1, 0], Kx=Kx, Ky=Ky, spectral=spectral) 264 | 265 | # If the data is in physical space, convert it to spectral space using 2D FFT. 266 | if spectral == False: 267 | Tau11_hat = np.fft.fft2(Tau11) 268 | Tau12_hat = np.fft.fft2(Tau12) 269 | Tau22_hat = np.fft.fft2(Tau22) 270 | U1x_hat = np.fft.fft2(Ux) 271 | U1y_hat = np.fft.fft2(Uy) 272 | V1x_hat = np.fft.fft2(Vx) 273 | else: # If the data is already in the spectral space, we don't need to transform it. 274 | Tau11_hat = Tau11 275 | Tau12_hat = Tau12 276 | Tau22_hat = Tau22 277 | U1x_hat = Ux 278 | U1y_hat = Uy 279 | V1x_hat = Vx 280 | 281 | # Assuming that U is a square matrix, derive the size of the LES grid. 282 | N_LES = U.shape[0] 283 | 284 | # Compute the energy transfer function in spectral space. 285 | Ptau_hat = (-np.conj(Tau11_hat)*U1x_hat + np.conj(Tau22_hat)*U1x_hat - np.conj(Tau12_hat)*U1y_hat - np.conj(Tau12_hat)*V1x_hat)/(N_LES*N_LES) 286 | 287 | # If the method is 'PiOmega', calculate the energy transfer using Psi and PiOmega. 288 | elif method == 'PiOmega': 289 | # Check if the necessary parameters are provided. 290 | if Psi is None or PiOmega is None: 291 | raise ValueError("Psi, PiOmega must be provided to calculate energy transfer using 'PiOmega' method") 292 | 293 | # If the data is in physical space, convert it to spectral space using 2D FFT. 294 | if spectral == False: 295 | Psi1_hat = np.fft.fft2(Psi) 296 | PiOmega_hat = np.fft.fft2(PiOmega) 297 | else: # If the data is already in the spectral space, we don't need to transform it. 298 | Psi1_hat = Psi 299 | PiOmega_hat = PiOmega 300 | 301 | # Assuming that Psi is a square matrix, derive the size of the LES grid. 302 | N_LES = Psi.shape[0] 303 | 304 | # Compute the energy transfer function in spectral space. 305 | Ptau_hat = (np.conj(PiOmega_hat)*Psi1_hat)/(N_LES*N_LES) 306 | 307 | else: # If an unsupported method is provided, raise an error. 308 | raise ValueError("Invalid method. Choose either 'Tau' or 'PiOmega'") 309 | 310 | # Perform an angle-averaged spectrum calculation on the computed energy transfer function. 311 | spectra, wavenumber = spectrum_angled_average(Ptau_hat, spectral=True) 312 | 313 | # Return the real part of the computed spectrum along with the corresponding wavenumbers. 314 | return np.real(spectra), wavenumber 315 | 316 | 317 | def enstrophyTransfer_spectra(Kx, Ky, Omega=None, Sigma1=None, Sigma2=None, PiOmega=None, method='Sigma', spectral=False): 318 | ''' 319 | Compute the enstrophy transfer spectra using 2D Forced Homogeneous Isotropic Turbulence (2D-FHIT) 320 | PZ > 0: dissipation 321 | PZ < 0: backscatter 322 | 323 | Parameters 324 | ---------- 325 | Kx, Ky : numpy.ndarray 326 | The 2D wavenumber arrays 327 | Omega : numpy.ndarray, optional 328 | The vorticity field (2D array) 329 | Sigma1 : numpy.ndarray, optional 330 | The Sigma1 field (2D array) 331 | Sigma2 : numpy.ndarray, optional 332 | The Sigma2 field (2D array) 333 | PiOmega : numpy.ndarray, optional 334 | The PiOmega field (2D array) 335 | method : str, optional 336 | The method to calculate Pz_hat. 'Sigma' (default) or 'PiOmega' 337 | spectral : bool, optional 338 | Determines whether the input is in the 'physical' space (default=False) or 'spectral' space 339 | 340 | Returns 341 | ------- 342 | spectra : numpy.ndarray 343 | The computed enstrophy transfer spectra 344 | wavenumber : numpy.ndarray 345 | The corresponding wavenumbers 346 | 347 | This function computes the enstrophy transfer spectra of the given 2D velocity fields 348 | using the 2D Forced Homogeneous Isotropic Turbulence. The inputs can be in the physical or spectral space. 349 | ''' 350 | 351 | # Assuming that Omega is a square matrix, the size of the LES grid is derived from Omega. 352 | N_LES = Omega.shape[0] 353 | 354 | # For the 'Sigma' method, compute the transfer spectra using Omega, Sigma1, and Sigma2. 355 | if method == 'Sigma': 356 | # Ensuring that the required arrays are provided. 357 | if Omega is None or Sigma1 is None or Sigma2 is None: 358 | raise ValueError("Omega, Sigma1, Sigma2 must be provided to calculate enstrophy transfer using 'Sigma' method") 359 | 360 | # Calculating the derivatives of Omega in x and y direction. 361 | Omegax = derivative(Omega, [1, 0], Kx=Kx, Ky=Ky, spectral=spectral) 362 | Omegay = derivative(Omega, [0, 1], Kx=Kx, Ky=Ky, spectral=spectral) 363 | 364 | # If the input is in the physical space, we need to convert it to spectral space using 2D FFT. 365 | if not spectral: 366 | Omegax_hat = np.fft.fft2(Omegax) 367 | Omegay_hat = np.fft.fft2(Omegay) 368 | Sigma1_hat = np.fft.fft2(Sigma1) 369 | Sigma2_hat = np.fft.fft2(Sigma2) 370 | else: # If the input is already in the spectral space, we don't need to transform it. 371 | Omegax_hat = Omegax 372 | Omegay_hat = Omegay 373 | Sigma1_hat = Sigma1 374 | Sigma2_hat = Sigma2 375 | 376 | # Computing the enstrophy transfer function in spectral space. 377 | Pz_hat = -(np.conj(Sigma1_hat) * Omegax_hat + np.conj(Sigma2_hat) * Omegay_hat) / (N_LES * N_LES) 378 | 379 | # For the 'PiOmega' method, compute the transfer spectra using Omega and PiOmega. 380 | elif method == 'PiOmega': 381 | # Ensuring that the required arrays are provided. 382 | if Omega is None or PiOmega is None: 383 | raise ValueError("Omega, PiOmega must be provided to calculate enstrophy transfer using 'PiOmega' method") 384 | 385 | # If the input is in the physical space, we need to convert it to spectral space using 2D FFT. 386 | if not spectral: 387 | Omega_hat = np.fft.fft2(Omega) 388 | PiOmega_hat = np.fft.fft2(PiOmega) 389 | else: # If the input is already in the spectral space, we don't need to transform it. 390 | Omega_hat = Omega 391 | PiOmega_hat = PiOmega 392 | 393 | # Computing the enstrophy transfer function in spectral space. Dividing by N_LES**2 is scale the fourier transform. 394 | Pz_hat = (np.conj(PiOmega_hat) * Omega_hat) / (N_LES * N_LES) 395 | 396 | else: # If an unsupported method is provided, raise an error. 397 | raise ValueError("Invalid method. Choose either 'Sigma' or 'PiOmega'") 398 | 399 | # Performing an angle-averaged spectrum calculation on the computed enstrophy transfer function. 400 | spectra, wavenumber = spectrum_angled_average(Pz_hat, spectral=True) 401 | 402 | # Returning the real part of the computed spectrum along with the corresponding wavenumbers. 403 | return np.real(spectra), wavenumber 404 | -------------------------------------------------------------------------------- /py2d/structure_function.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def structure_fxn_scalar(data, orderMax=8, Lx=2*np.pi, Ly=2*np.pi): 4 | """Calculates the scalar velocity structure function up to a specified order. 5 | 6 | Args: 7 | data (np.ndarray): A 2D array representing the scalar field. 8 | orderMax (int, optional): The maximum order of the structure function to calculate. 9 | Defaults to 8. 10 | Lx (float, optional): Physical size of the domain in the x-direction. Defaults to 2*pi. 11 | Ly (float, optional): Physical size of the domain in the y-direction. Defaults to 2*pi. 12 | 13 | Returns: 14 | tuple: A tuple containing: 15 | * np.ndarray: A 3D array containing the structure function values. 16 | The dimensions are (lx, ly, order). 17 | * np.ndarray: Array of physical separation distances (lx, ly) 18 | The dimensions are (lx, ly, order). 19 | 20 | Note: 21 | Reference: https://joss.theoj.org/papers/10.21105/joss.02185 22 | """ 23 | 24 | Nx, Ny = data.shape 25 | dx = Lx/Nx # Grid spacing in x-direction 26 | dy = Ly/Ny # Grid spacing in y-direction 27 | 28 | # Initialize the structure function aarrays 29 | structure_fxn = np.zeros((Nx//2, Ny//2, orderMax)) 30 | l_array = np.zeros((Nx//2, Ny//2, 2)) 31 | 32 | # Iterate over separation distances (indices) 33 | for lx_ind in range(1, Nx//2 + 1): 34 | for ly_ind in range(1, Ny//2 + 1): 35 | ds = cal_structure_fxn_diff(data, lx_ind, ly_ind) 36 | 37 | # Calculate structure functions for different orders 38 | for order in range(1, orderMax + 1): 39 | structure_fxn[lx_ind-1, ly_ind-1, order-1] = np.mean(ds**order) # Zero-based indexing 40 | 41 | l_array[lx_ind-1, ly_ind-1, 0] = lx_ind*dx 42 | l_array[lx_ind-1, ly_ind-1, 1] = ly_ind*dy 43 | 44 | return structure_fxn, l_array 45 | 46 | def structure_fxn_2Dvector(dataA, dataB, orderMax=8, Lx=2*np.pi, Ly=2*np.pi): 47 | """Calculates longitudinal and transverse velocity structure functions for a 2D vector field. 48 | 49 | Args: 50 | dataA (np.ndarray): A 2D array representing the first component of the vector field. 51 | dataB (np.ndarray): A 2D array representing the second component of the vector field. 52 | orderMax (int, optional): The maximum order of the structure function to calculate. 53 | Defaults to 8. 54 | Lx (float, optional): Physical size of the domain in the x-direction. Defaults to 2*pi. 55 | Ly (float, optional): Physical size of the domain in the y-direction. Defaults to 2*pi. 56 | 57 | Returns: 58 | tuple: A tuple containing: 59 | * np.ndarray: Longitudinal velocity structure function (3D array) 60 | * np.ndarray: Transverse velocity structure function (3D array) 61 | * np.ndarray: Array of physical separation distances (lx, ly, 2) 62 | 63 | Note: 64 | Reference: https://joss.theoj.org/papers/10.21105/joss.02185 65 | """ 66 | 67 | Nx, Ny = dataA.shape 68 | 69 | dx = Lx/Nx # Grid spacing in x-direction 70 | dy = Ly/Ny # Grid spacing in y-direction 71 | 72 | # Initialize the structure function arrays 73 | structure_fxn_longitudinal = np.zeros((Nx//2, Ny//2, orderMax)) 74 | structure_fxn_transverse = np.zeros((Nx//2, Ny//2, orderMax)) 75 | l_array = np.zeros((Nx//2, Ny//2, 2)) 76 | 77 | # Iterate over separation distances (indices) 78 | for lx_ind in range(1, Nx//2 + 1): 79 | for ly_ind in range(1, Ny//2 + 1): 80 | du = cal_structure_fxn_diff(dataA, lx_ind, ly_ind) 81 | dv = cal_structure_fxn_diff(dataB, lx_ind, ly_ind) 82 | 83 | # Physical separation distances 84 | lx = lx_ind * dx 85 | ly = ly_ind * dy 86 | 87 | # Unit vector calculation 88 | l_norm = np.sqrt(lx**2 + ly**2) 89 | lx_unit = lx / l_norm 90 | ly_unit = ly / l_norm 91 | 92 | # Projection along and perpendicular to separation vector 93 | diff_parallel = du*lx_unit + dv*ly_unit 94 | diff_perp = np.sqrt((du - diff_parallel*lx_unit)**2 + (dv - diff_parallel*ly_unit)**2) 95 | 96 | # Calculate structure functions for different orders 97 | for order in range(1, orderMax + 1): 98 | structure_fxn_longitudinal[lx_ind-1, ly_ind-1, order-1] = np.mean(diff_parallel**order) 99 | structure_fxn_transverse[lx_ind-1, ly_ind-1, order-1] = np.mean(diff_perp**order) 100 | 101 | l_array[lx_ind-1, ly_ind-1, 0] = lx 102 | l_array[lx_ind-1, ly_ind-1, 1] = ly 103 | 104 | return structure_fxn_longitudinal, structure_fxn_transverse, l_array 105 | 106 | def cal_structure_fxn_diff(data, lx_ind, ly_ind): 107 | """Calculates the difference for the velocity structure function calculation. 108 | 109 | Args: 110 | data (np.ndarray): The 2D array representing the field. 111 | lx (int): The separation distance in the x-direction. 112 | ly (int): The separation distance in the y-direction. 113 | Nx (int): The number of grid points in the x-direction. 114 | Ny (int): The number of grid points in the y-direction. 115 | 116 | Returns: 117 | np.ndarray: The difference array used for calculating the structure function. 118 | """ 119 | 120 | Nx, Ny = data.shape 121 | 122 | # Calculate the velocity difference u(x+r) - u(x). 123 | data_diff = data[lx_ind:Nx//2 + lx_ind, ly_ind:Ny//2 + ly_ind] - data[:Nx//2, :Ny//2] 124 | 125 | return data_diff 126 | 127 | def angled_average_2D(A, l_array, kmax = 'grid'): 128 | '''Calculates the angle-averaged 2D scalar matrix''' 129 | 130 | # Check if input 'A' is a 2D square matrix 131 | if not isinstance(A, np.ndarray) or A.ndim != 2 or A.shape[0] != A.shape[1]: 132 | raise ValueError('Input is not a 2D square matrix. Please input a 2D square matrix') 133 | # Check if input 'A' contains non-numeric values 134 | if not np.issubdtype(A.dtype, np.number): 135 | raise ValueError('Input contains non-numeric values') 136 | 137 | lx = l_array[:,:,0] 138 | ly = l_array[:,:,1] 139 | labs = np.sqrt(lx**2 + ly**2) 140 | dlx = lx[1,0] - lx[0,0] 141 | 142 | # Calculate the maximum wavenumber to be considered in the average 143 | if kmax == 'grid': 144 | l = np.arange(np.min(lx), np.max(lx)+dlx, dlx) 145 | elif kmax == 'diagonal': 146 | l = np.arange(np.min(labs), np.max(labs), dlx) 147 | 148 | # Initialize the output array with zeros 149 | A_angled_average = np.zeros(np.shape(l)) 150 | 151 | # Compute the angle-averaged 152 | for l_ind, l_value in enumerate(l): 153 | l_tempInd = (labs > (l_value - dlx/2)) & (labs <= (l_value + dlx/2)) 154 | A_angled_average[l_ind] = np.sum(A[l_tempInd]) 155 | 156 | return A_angled_average, l 157 | 158 | 159 | def angled_average_struc_fxn_2D(structure_fx, l_array, kmax='grid'): 160 | """Calculating the spectra of the whole structure function or flatness containing structure functions upto orderMax""" 161 | 162 | orderMax = structure_fx.shape[2] 163 | 164 | angled_average_structure_fxn_list = [] 165 | for count in range(orderMax): 166 | structure_fxn_temp, l_struc_fxn = angled_average_2D(structure_fx[:,:,count], l_array, kmax=kmax) 167 | angled_average_structure_fxn_list.append(structure_fxn_temp) 168 | 169 | angled_average_structure_fxn = np.array(angled_average_structure_fxn_list).T 170 | 171 | return angled_average_structure_fxn, l_struc_fxn 172 | 173 | 174 | 175 | def flatness(structure_fxn_arr): 176 | """Calculates the flatness of the structure function. The structure function array should be output of 177 | 178 | Args: 179 | structure_fxn (np.ndarray): The structure function values 180 | orderMax (int, optional): The maximum order of the structure function. Defaults to 8. 181 | 182 | Returns: 183 | np.ndarray: The flatness values. 184 | """ 185 | 186 | flatness_arr = np.zeros(structure_fxn_arr.shape) 187 | for order in range(1,structure_fxn_arr.shape[2]+1): 188 | flatness_arr[:,:,order-1] = structure_fxn_arr[:,:,order-1] ** order/ (structure_fxn_arr[:,:,1]**(order/2)) 189 | 190 | return flatness_arr -------------------------------------------------------------------------------- /py2d/util.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import jax.numpy as jnp 3 | from scipy.interpolate import RegularGridInterpolator 4 | 5 | def corr2(a_var,b_var): 6 | # Correlation coefficient of N x N x T array 7 | 8 | a = a_var - np.mean(a_var) 9 | b = b_var - np.mean(b_var) 10 | 11 | r = (a*b).sum() / np.sqrt((a*a).sum() * (b*b).sum()) 12 | 13 | return r 14 | 15 | # Function to convert fft2 outputs to rfft2 outputs 16 | def fft2_to_rfft2(a_hat_fft): 17 | if a_hat_fft.shape[0] % 2 == 0: 18 | # Shape of matrix is even 19 | return a_hat_fft[:,:a_hat_fft.shape[1]//2+1] 20 | else: 21 | # Shape of matrix is odd 22 | return a_hat_fft[:,:(a_hat_fft.shape[1]-1)//2+1] 23 | 24 | # Function to convert rfft2 outputs to fft2 outputs 25 | def rfft2_to_fft2(a_hat_rfft2, backend=np): 26 | 27 | a_hat_fft2 = backend.zeros([a_hat_rfft2.shape[0], a_hat_rfft2.shape[0]], dtype=backend.complex128) 28 | 29 | if a_hat_rfft2.shape[0] % 2 == 0: 30 | # Shape of matrix is even 31 | kn = a_hat_rfft2.shape[0]//2 # Index of Nyquist wavenumber 32 | a_hat_fft2[:,:kn+1] = a_hat_rfft2 33 | a_hat_fft2[0,kn+1:] = backend.flip(backend.conj(a_hat_rfft2[0,1:kn])) # Making the first wavenumber conjugate symmetric 34 | a_hat_fft2[1:,kn+1:] = backend.flip(backend.flip(backend.conj(a_hat_rfft2[1:,1:kn]),axis=0),axis=1) # Making the rest of the matrix conjugate symmetric 35 | 36 | else: 37 | # Shape of matrix is odd 38 | kn = (a_hat_rfft2.shape[0]-1)//2 39 | a_hat_fft2[:,:kn+1] = a_hat_rfft2 40 | a_hat_fft2[0,kn+1:] = backend.flip(backend.conj(a_hat_rfft2[0,1:kn+1])) # Making the first wavenumber conjugate symmetric 41 | a_hat_fft2[1:,kn+1:] = backend.flip(backend.flip(backend.conj(a_hat_rfft2[1:,1:kn+1]),axis=0),axis=1) # Making the rest of the matrix conjugate symmetric 42 | 43 | return a_hat_fft2 44 | 45 | def regrid(data, out_x, out_y): 46 | # Regrid data to a new grid size 47 | m = max(data.shape[0], data.shape[1]) 48 | y = np.linspace(0, 1.0/m, data.shape[0]) 49 | x = np.linspace(0, 1.0/m, data.shape[1]) 50 | interpolating_function = RegularGridInterpolator((y, x), data) 51 | 52 | yv, xv = np.meshgrid(np.linspace(0, 1.0/m, out_y), np.linspace(0, 1.0/m, out_x)) 53 | 54 | return interpolating_function((xv, yv)) 55 | 56 | def eig_vec_2D(A11, A12, A21, A22): 57 | """ 58 | Calculate eigenvectors and eigenvalues of a 2D tensor field. 59 | 60 | This function takes three 2D arrays representing the components of a 2D tensor field and computes 61 | the eigenvalues and eigenvectors for each 2x2 tensor at every point in the field. 62 | 63 | Parameters: 64 | A11 (np.ndarray): A 2D array representing the A11 component of the tensor field. 65 | A12 (np.ndarray): A 2D array representing the A12 component of the tensor field. 66 | A22 (np.ndarray): A 2D array representing the A22 component of the tensor field. 67 | 68 | Returns: 69 | tuple: A tuple containing the following elements: 70 | - eigVec1 (np.ndarray): A 2xN array where N is the number of elements in A11. Each column represents the first eigenvector of the corresponding tensor. 71 | - eigVec2 (np.ndarray): A 2xN array where N is the number of elements in A11. Each column represents the second eigenvector of the corresponding tensor. 72 | - eigVal1 (np.ndarray): A 1D array of length N, containing the first eigenvalue of each tensor. 73 | - eigVal2 (np.ndarray): A 1D array of length N, containing the second eigenvalue of each tensor. 74 | """ 75 | # Initialize eigenvalues and eigenvectors 76 | eigVal1 = np.zeros(A11.size) 77 | eigVal2 = np.zeros(A11.size) 78 | eigVec1 = np.zeros((A11.size,2)) 79 | eigVec2 = np.zeros((A11.size,2)) 80 | 81 | # Loop through each element 82 | for countGrid in range(A11.size): 83 | A = np.array([[A11.flat[countGrid], A12.flat[countGrid]], 84 | [A12.flat[countGrid], A22.flat[countGrid]]]) 85 | eigVals, eigVecs = np.linalg.eig(A) 86 | # Sort the eigenvalues and eigenvectors 87 | idx = eigVals.argsort()[::-1] 88 | eigVals = eigVals[idx] 89 | eigVecs = eigVecs[:,idx] 90 | 91 | eigVec1[countGrid,:] = eigVecs[:, 0] 92 | eigVec2[countGrid,:] = eigVecs[:, 1] 93 | eigVal1[countGrid] = eigVals[0] 94 | eigVal2[countGrid] = eigVals[1] 95 | 96 | return eigVec1, eigVec2, eigVal1, eigVal2 97 | 98 | 99 | -------------------------------------------------------------------------------- /py2d/uv2tau_CNN.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import jax.numpy as jnp 5 | 6 | # Function to initialize the model 7 | def init_model(model_type, model_path): 8 | # Define the ConvNeuralNet classes for each model type 9 | if model_type == '1M': 10 | class ConvNeuralNet(nn.Module): 11 | 12 | # Determine what layers and their order in CNN object 13 | def __init__(self, num_classes): 14 | 15 | super(ConvNeuralNet, self).__init__() 16 | 17 | self.conv_layer1 = nn.Conv2d(in_channels=2, out_channels=64, kernel_size=5, padding="same") 18 | self.conv_layer2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 19 | self.conv_layer3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 20 | self.conv_layer4 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 21 | self.conv_layer5 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 22 | self.conv_layer6 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 23 | self.conv_layer7 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 24 | self.conv_layer8 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 25 | self.conv_layer9 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 26 | self.conv_layer10 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 27 | self.conv_layer11 = nn.Conv2d(in_channels=64, out_channels=3, kernel_size=5, padding="same") 28 | self.relu1 = nn.ReLU() 29 | 30 | def forward(self, x): 31 | 32 | out = self.conv_layer1(x) # Layer1 (Input Layer) 33 | out = self.relu1(out) 34 | 35 | ## Hidden Layers 36 | out = self.conv_layer2(out) #Layer2 37 | out = self.relu1(out) 38 | 39 | out = self.conv_layer3(out) #Layer3 40 | out = self.relu1(out) 41 | 42 | out = self.conv_layer4(out) #Layer4 43 | out = self.relu1(out) 44 | 45 | out = self.conv_layer5(out) #Layer5 46 | out = self.relu1(out) 47 | 48 | out = self.conv_layer6(out) #Layer6 49 | out = self.relu1(out) 50 | 51 | out = self.conv_layer7(out) #Layer7 52 | out = self.relu1(out) 53 | 54 | out = self.conv_layer8(out) #Layer8 55 | out = self.relu1(out) 56 | 57 | out = self.conv_layer9(out) #Layer9 58 | out = self.relu1(out) 59 | 60 | out = self.conv_layer10(out) #Layer10 61 | out = self.relu1(out) 62 | 63 | out = self.conv_layer11(out) #Layer11 (Output Layer) 64 | return out 65 | elif model_type == '1M-PiOmega': 66 | class ConvNeuralNet(nn.Module): 67 | 68 | # Determine what layers and their order in CNN object 69 | def __init__(self, num_classes): 70 | 71 | super(ConvNeuralNet, self).__init__() 72 | 73 | self.conv_layer1 = nn.Conv2d(in_channels=2, out_channels=64, kernel_size=5, padding="same") 74 | self.conv_layer2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 75 | self.conv_layer3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 76 | self.conv_layer4 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 77 | self.conv_layer5 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 78 | self.conv_layer6 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 79 | self.conv_layer7 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 80 | self.conv_layer8 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 81 | self.conv_layer9 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 82 | self.conv_layer10 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 83 | self.conv_layer11 = nn.Conv2d(in_channels=64, out_channels=1, kernel_size=5, padding="same") 84 | self.relu1 = nn.ReLU() 85 | 86 | def forward(self, x): 87 | 88 | out = self.conv_layer1(x) # Layer1 (Input Layer) 89 | out = self.relu1(out) 90 | 91 | ## Hidden Layers 92 | out = self.conv_layer2(out) #Layer2 93 | out = self.relu1(out) 94 | 95 | out = self.conv_layer3(out) #Layer3 96 | out = self.relu1(out) 97 | 98 | out = self.conv_layer4(out) #Layer4 99 | out = self.relu1(out) 100 | 101 | out = self.conv_layer5(out) #Layer5 102 | out = self.relu1(out) 103 | 104 | out = self.conv_layer6(out) #Layer6 105 | out = self.relu1(out) 106 | 107 | out = self.conv_layer7(out) #Layer7 108 | out = self.relu1(out) 109 | 110 | out = self.conv_layer8(out) #Layer8 111 | out = self.relu1(out) 112 | 113 | out = self.conv_layer9(out) #Layer9 114 | out = self.relu1(out) 115 | 116 | out = self.conv_layer10(out) #Layer10 117 | out = self.relu1(out) 118 | 119 | out = self.conv_layer11(out) #Layer11 (Output Layer) 120 | return out 121 | elif model_type == 'mcwiliams-ani': 122 | class ConvNeuralNet(nn.Module): 123 | 124 | # Determine what layers and their order in CNN object 125 | def __init__(self, num_classes): 126 | 127 | super(ConvNeuralNet, self).__init__() 128 | 129 | self.conv_layer1 = nn.Conv2d(in_channels=2, out_channels=64, kernel_size=5, padding="same") 130 | self.conv_layer2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 131 | self.conv_layer3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 132 | self.conv_layer4 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 133 | self.conv_layer5 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 134 | self.conv_layer6 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 135 | self.conv_layer7 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 136 | self.conv_layer8 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 137 | self.conv_layer9 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 138 | self.conv_layer10 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 139 | self.conv_layer11 = nn.Conv2d(in_channels=64, out_channels=2, kernel_size=5, padding="same") 140 | self.relu1 = nn.ReLU() 141 | 142 | def forward(self, x): 143 | 144 | out = self.conv_layer1(x) # Layer1 (Input Layer) 145 | out = self.relu1(out) 146 | 147 | ## Hidden Layers 148 | out = self.conv_layer2(out) #Layer2 149 | out = self.relu1(out) 150 | 151 | out = self.conv_layer3(out) #Layer3 152 | out = self.relu1(out) 153 | 154 | out = self.conv_layer4(out) #Layer4 155 | out = self.relu1(out) 156 | 157 | out = self.conv_layer5(out) #Layer5 158 | out = self.relu1(out) 159 | 160 | out = self.conv_layer6(out) #Layer6 161 | out = self.relu1(out) 162 | 163 | out = self.conv_layer7(out) #Layer7 164 | out = self.relu1(out) 165 | 166 | out = self.conv_layer8(out) #Layer8 167 | out = self.relu1(out) 168 | 169 | out = self.conv_layer9(out) #Layer9 170 | out = self.relu1(out) 171 | 172 | out = self.conv_layer10(out) #Layer10 173 | out = self.relu1(out) 174 | 175 | out = self.conv_layer11(out) #Layer11 (Output Layer) 176 | return out 177 | 178 | elif model_type == 'mcwiliams': 179 | class ConvNeuralNet(nn.Module): 180 | 181 | # Determine what layers and their order in CNN object 182 | def __init__(self, num_classes): 183 | 184 | super(ConvNeuralNet, self).__init__() 185 | 186 | self.conv_layer1 = nn.Conv2d(in_channels=2, out_channels=32, kernel_size=5, padding="same") 187 | # self.conv_layer2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 188 | # self.conv_layer3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 189 | # self.conv_layer4 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 190 | # self.conv_layer5 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 191 | # self.conv_layer6 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 192 | # self.conv_layer7 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 193 | # self.conv_layer8 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 194 | # self.conv_layer9 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 195 | # self.conv_layer10 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 196 | self.conv_layer11 = nn.Conv2d(in_channels=32, out_channels=3, kernel_size=5, padding="same") 197 | self.relu1 = nn.ReLU() 198 | 199 | def forward(self, x): 200 | 201 | out = self.conv_layer1(x) # Layer1 (Input Layer) 202 | out = self.relu1(out) 203 | 204 | ## Hidden Layers 205 | # out = self.conv_layer2(out) #Layer2 206 | # out = self.relu1(out) 207 | 208 | # out = self.conv_layer3(out) #Layer3 209 | # out = self.relu1(out) 210 | 211 | # out = self.conv_layer4(out) #Layer4 212 | # out = self.relu1(out) 213 | 214 | # out = self.conv_layer5(out) #Layer5 215 | # out = self.relu1(out) 216 | 217 | # out = self.conv_layer6(out) #Layer6 218 | # out = self.relu1(out) 219 | 220 | # out = self.conv_layer7(out) #Layer7 221 | # out = self.relu1(out) 222 | 223 | # out = self.conv_layer8(out) #Layer8 224 | # out = self.relu1(out) 225 | 226 | # out = self.conv_layer9(out) #Layer9 227 | # out = self.relu1(out) 228 | 229 | # out = self.conv_layer10(out) #Layer10 230 | # out = self.relu1(out) 231 | 232 | out = self.conv_layer11(out) #Layer11 (Output Layer) 233 | return out 234 | 235 | elif model_type == '8M': 236 | class ConvNeuralNet(nn.Module): 237 | # Determine what layers and their order in CNN object 238 | def __init__(self, num_classes): 239 | 240 | super(ConvNeuralNet, self).__init__() 241 | 242 | self.conv_layer1 = nn.Conv2d(in_channels=2, out_channels=64, kernel_size=15, padding="same") 243 | self.conv_layer2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=15, padding="same") 244 | self.conv_layer3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=15, padding="same") 245 | self.conv_layer4 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=15, padding="same") 246 | self.conv_layer5 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=15, padding="same") 247 | self.conv_layer6 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=15, padding="same") 248 | self.conv_layer7 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=15, padding="same") 249 | self.conv_layer8 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=15, padding="same") 250 | self.conv_layer9 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=15, padding="same") 251 | self.conv_layer10 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=15, padding="same") 252 | self.conv_layer11 = nn.Conv2d(in_channels=64, out_channels=3, kernel_size=15, padding="same") 253 | self.relu1 = nn.ReLU() 254 | 255 | def forward(self, x): 256 | 257 | out = self.conv_layer1(x) # Layer1 (Input Layer) 258 | out = self.relu1(out) 259 | 260 | ## Hidden Layers 261 | out = self.conv_layer2(out) #Layer2 262 | out = self.relu1(out) 263 | 264 | out = self.conv_layer3(out) #Layer3 265 | out = self.relu1(out) 266 | 267 | out = self.conv_layer4(out) #Layer4 268 | out = self.relu1(out) 269 | 270 | out = self.conv_layer5(out) #Layer5 271 | out = self.relu1(out) 272 | 273 | out = self.conv_layer6(out) #Layer6 274 | out = self.relu1(out) 275 | 276 | out = self.conv_layer7(out) #Layer7 277 | out = self.relu1(out) 278 | 279 | out = self.conv_layer8(out) #Layer8 280 | out = self.relu1(out) 281 | 282 | out = self.conv_layer9(out) #Layer9 283 | out = self.relu1(out) 284 | 285 | out = self.conv_layer10(out) #Layer10 286 | out = self.relu1(out) 287 | 288 | out = self.conv_layer11(out) #Layer11 (Output Layer) 289 | return out 290 | 291 | elif model_type == 'shallow': 292 | class ConvNeuralNet(nn.Module): 293 | # Determine what layers and their order in CNN object 294 | def __init__(self, num_classes): 295 | 296 | super(ConvNeuralNet, self).__init__() 297 | 298 | self.conv_layer1 = nn.Conv2d(in_channels=2, out_channels=64, kernel_size=5, padding="same") 299 | self.conv_layer2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 300 | self.conv_layer3 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 301 | self.conv_layer4 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 302 | self.conv_layer5 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 303 | self.conv_layer6 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 304 | self.conv_layer7 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 305 | self.conv_layer8 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 306 | self.conv_layer9 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 307 | self.conv_layer10 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, padding="same") 308 | self.conv_layer11 = nn.Conv2d(in_channels=64, out_channels=3, kernel_size=5, padding="same") 309 | self.relu1 = nn.ReLU() 310 | 311 | def forward(self, x): 312 | 313 | out = self.conv_layer1(x) # Layer1 (Input Layer) 314 | out = self.relu1(out) 315 | 316 | ## Hidden Layers 317 | out = self.conv_layer2(out) #Layer2 318 | out = self.relu1(out) 319 | 320 | # out = self.conv_layer3(out) #Layer3 321 | # out = self.relu1(out) 322 | 323 | # out = self.conv_layer4(out) #Layer4 324 | # out = self.relu1(out) 325 | 326 | # out = self.conv_layer5(out) #Layer5 327 | # out = self.relu1(out) 328 | 329 | # out = self.conv_layer6(out) #Layer6 330 | # out = self.relu1(out) 331 | 332 | # out = self.conv_layer7(out) #Layer7 333 | # out = self.relu1(out) 334 | 335 | # out = self.conv_layer8(out) #Layer8 336 | # out = self.relu1(out) 337 | 338 | # out = self.conv_layer9(out) #Layer9 339 | # out = self.relu1(out) 340 | 341 | # out = self.conv_layer10(out) #Layer10 342 | # out = self.relu1(out) 343 | 344 | out = self.conv_layer11(out) #Layer11 (Output Layer) 345 | return out 346 | else: 347 | print('Invalid model type. Please choose from "deep" or "shallow"') 348 | 349 | # Set device to GPU if available, otherwise use CPU 350 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 351 | 352 | # Create model instance and move it to the specified device 353 | model = ConvNeuralNet(1).to(device=device) 354 | 355 | # Load the checkpoint 356 | checkpoint = torch.load(model_path, map_location=device) 357 | 358 | # Extract the model_state_dict from the checkpoint 359 | model_state_dict = checkpoint['model_state_dict'] 360 | 361 | # Load the state dictionary into the model 362 | model.load_state_dict(model_state_dict) 363 | 364 | # Move the model to the specified device 365 | model.to(device) 366 | 367 | # Set the model to evaluation mode 368 | model.eval() 369 | 370 | return model 371 | 372 | 373 | def evaluate_model(model, input_tensor): 374 | NX = input_tensor.shape[1] 375 | NY = input_tensor.shape[2] 376 | # Set device to GPU if available, otherwise use CPU 377 | # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 378 | 379 | # Convert the JAX array to a NumPy array, then to a PyTorch tensor 380 | # input_np = np.array(input_tensor) 381 | # input_tensor = torch.from_numpy(input_tensor).float() 382 | input_tensor = j2t(input_tensor).float() 383 | 384 | # Convert input_tensor to a torch Variable and move it to the specified device 385 | # input_var = torch.autograd.Variable(input_tensor).to(device) 386 | 387 | # Perform forward pass of the model with the input variable 388 | # output = model(input_var) 389 | output = model(input_tensor) 390 | # print('Output shape right after model: ', output.shape) 391 | # Convert the output tensor back to a JAX array 392 | # output_np = output.detach().cpu().numpy() 393 | # output_jax = jnp.array(output_np) 394 | # output = torch.concatenate([output, output, output], dim=0) 395 | # Flatten the output tensor 396 | output = output.view(-1) 397 | output_jax = t2j(output) 398 | # Reshape the output tensor to the original shape 399 | output_jax = output_jax.reshape((-1,NX,NY)) 400 | 401 | 402 | return output_jax 403 | 404 | import torch 405 | import torch.utils.dlpack 406 | import jax 407 | import jax.dlpack 408 | import jax.numpy as np 409 | import numpy as nnp 410 | import time 411 | 412 | # A generic mechanism for turning a JAX function into a PyTorch function. 413 | 414 | def j2t(x_jax): 415 | x_torch = torch.utils.dlpack.from_dlpack(jax.dlpack.to_dlpack(x_jax)) 416 | return x_torch 417 | 418 | def t2j(x_torch): 419 | x_torch = x_torch.contiguous() # https://github.com/google/jax/issues/8082 420 | x_jax = jax.dlpack.from_dlpack(torch.utils.dlpack.to_dlpack(x_torch)) 421 | return x_jax 422 | 423 | 424 | 425 | ''' 426 | ------------------------------------------- Examples of how to use the functions above ------------------------------------------- 427 | model_path = "/media/volume/sdb/cgft/cgft_shallow/train_model_0b898_00008_8_batch_size_train=32,learning_rate=0.0000,p_data=50_2023-04-17_23-57-20/model.pt" 428 | input_tensor = torch.randn(5, 2, 128, 128) # Example input tensor, modify according to your model input requirements 429 | 430 | model = init_model(model_type='1M', model_path=model_path) 431 | output = evaluate_model(model, input_tensor) 432 | 433 | print(output.shape) 434 | ''' 435 | 436 | 437 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools >= 61"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "Py2D" 7 | version = "0.1" 8 | description = "Python solver for incompressible 2D Navier-Stokes equations." 9 | authors = [ 10 | { name="Karan Jakhar", email="karanj@uchicago.edu" }, 11 | ] 12 | readme = "README.md" 13 | license = {file = "LICENSE"} 14 | requires-python = ">=3.10" 15 | classifiers = [ 16 | "License :: OSI Approved :: MIT License", 17 | "Development Status :: 3 - Alpha", 18 | "Natural Language :: English", 19 | "Programming Language :: Python :: 3", 20 | 'Programming Language :: Python :: 3.10', 21 | 'Programming Language :: Python :: 3.11', 22 | 'Programming Language :: Python :: 3.12', 23 | "Operating System :: OS Independent", 24 | ] 25 | dependencies = [ 26 | "numpy", 27 | "jax", 28 | "jaxlib", 29 | ] 30 | 31 | [project.urls] 32 | "Homepage" = "https://github.com/envfluids/py2d" 33 | "Bug Tracker" = "https://github.com/envfluids/py2d/issues" 34 | 35 | [tool.setuptools] 36 | packages = ["py2d"] 37 | include-package-data = true --------------------------------------------------------------------------------