Using pyFC to generate an xmdf random velocity field with lognormal one-point statistics and Kolmogorov two-point correlation

fractal-cube.py

# -*- coding: utf-8 -*-
"""
Created on Tue Sep 10 16:33:49 2024

@author: alankar
"""

import numpy as np
import matplotlib.pyplot as plt
from pyfc import pyFC
import h5py

# TODO: fix data dump as per precision
ni, nj, nk = 128, 128, 128
prec = 'double'
xmin, xmax = 0., 1.
ymin, ymax = 0., 1.
zmin, zmax = 0., 1.
time = 0

kmin = 1./np.max([xmax, ymax, zmax])
kmax = np.max([(ni+1)/2., (nj+1)/2., (nk+1)/2.]) # nyquist limit

x, y, z = np.meshgrid(np.linspace(xmin, xmax, ni+1), 
                      np.linspace(ymin, ymax, nj+1),
                      np.linspace(zmin, zmax, nk+1))

fc = pyFC.LogNormalFractalCube(ni=ni, nj=nj, nk=nk, 
                               kmin=kmin, mean=1., 
                               sigma=np.sqrt(5.), beta=-5./3.)
cube = fc.gen_cube()

plt.ion()
# pyFC.plot_midplane_slice(), pyFC.plot_raytrace(), pyFC.plot_power_spec(), pyFC.plot_pdf()
pyFC.plot_midplane_slice(fc, ax=2, scaling='log', 
                         cmap='viridis', plottype='imshow', 
                         labels=True, colorbar=True, 
                         kminlabel=True, 
                         vmin=-2.1, vmax=2.1)

pyFC.plot_field_stats(fc, scaling='log', 
                      vmin=-2.1, vmax=2.1, 
                      cmap='viridis')

data = cube.cube
# pyFC.write_cube(fc, 'raw_field.bin')

vx, vy, vz = np.zeros_like(data), np.zeros_like(data), np.zeros_like(data)

for k in range(nk):
    for j in range(nj):
        for i in range(ni):
            vx_tmp, vy_tmp = np.random.uniform(low=0, high=data[k,j,i], size=(2,))
            while np.sqrt(vx_tmp**2 + vy_tmp**2)>data[k,j,i]:
                vx_tmp, vy_tmp = np.random.uniform(low=0, high=data[k,j,i], size=(2,))
            vz_tmp = np.sqrt(data[k,j,i]**2 - (vx_tmp**2 + vy_tmp**2) )
            vx[k,j,i] = np.random.choice([-1,1], 1, p=[0.5,0.5])[0] * vx_tmp
            vy[k,j,i] = np.random.choice([-1,1], 1, p=[0.5,0.5])[0] * vy_tmp
            vz[k,j,i] = np.random.choice([-1,1], 1, p=[0.5,0.5])[0] * vz_tmp

with h5py.File(f'data.{time:04d}.{"flt" if prec=="single" else "dbl"}.h5', 'w') as hdf:
    hdf.create_dataset('/node_coords/X', shape = x.T.shape, data = x.T)
    hdf.create_dataset('/node_coords/Y', shape = y.T.shape, data = y.T)
    hdf.create_dataset('/node_coords/Z', shape = z.T.shape, data = z.T)
    hdf.create_dataset(f'/Timestep_{time}/vars/VEL',
                       shape=data.T.shape, data = data.T)
    hdf.create_dataset(f'/Timestep_{time}/vars/VX1',
                       shape=vx.T.shape, data = vx.T)
    hdf.create_dataset(f'/Timestep_{time}/vars/VX2',
                       shape=vy.T.shape, data = vy.T)
    hdf.create_dataset(f'/Timestep_{time}/vars/VX3',
                       shape=vz.T.shape, data = vz.T)
    

xmf = f'''
<?xml version="1.0" ?>
<!DOCTYPE Xdmf SYSTEM "Xdmf.dtd" []>
<Xdmf Version="2.0">
 <Domain>
   <Grid Name="node_mesh" GridType="Uniform">
    <Time Value="{time}"/>
     <Topology TopologyType="3DSMesh" NumberOfElements="{ni+1} {nj+1} {nk+1}"/>
     <Geometry GeometryType="X_Y_Z">
       <DataItem Dimensions="{ni+1} {nj+1} {nk+1}" NumberType="Float" Precision="{"4" if prec=="single" else "8"}" Format="HDF">
        data.{time:04d}.{"flt" if prec=="single" else "dbl"}.h5:/node_coords/X
       </DataItem>
       <DataItem Dimensions="{ni+1} {nj+1} {nk+1}" NumberType="Float" Precision="{"4" if prec=="single" else "8"}" Format="HDF">
        data.{time:04d}.{"flt" if prec=="single" else "dbl"}.h5:/node_coords/Y
       </DataItem>
       <DataItem Dimensions="{ni+1} {nj+1} {nk+1}" NumberType="Float" Precision="{"4" if prec=="single" else "8"}" Format="HDF">
        data.{time:04d}.{"flt" if prec=="single" else "dbl"}.h5:/node_coords/Z
       </DataItem>
     </Geometry>
     <Attribute Name="vel" AttributeType="Scalar" Center="Cell">
       <DataItem Dimensions="{ni} {nj} {nk}" NumberType="Float" Precision="{"4" if prec=="single" else "8"}" Format="HDF">
        data.{time:04d}.{"flt" if prec=="single" else "dbl"}.h5:/Timestep_{time}/vars/VEL
       </DataItem>
     </Attribute>
     <Attribute Name="vx1" AttributeType="Scalar" Center="Cell">
       <DataItem Dimensions="{ni} {nj} {nk}" NumberType="Float" Precision="{"4" if prec=="single" else "8"}" Format="HDF">
        data.{time:04d}.{"flt" if prec=="single" else "dbl"}.h5:/Timestep_{time}/vars/VX1
       </DataItem>
     </Attribute>
     <Attribute Name="vx2" AttributeType="Scalar" Center="Cell">
       <DataItem Dimensions="{ni} {nj} {nk}" NumberType="Float" Precision="{"4" if prec=="single" else "8"}" Format="HDF">
        data.{time:04d}.{"flt" if prec=="single" else "dbl"}.h5:/Timestep_{time}/vars/VX2
       </DataItem>
     </Attribute>
     <Attribute Name="vx3" AttributeType="Scalar" Center="Cell">
       <DataItem Dimensions="{ni} {nj} {nk}" NumberType="Float" Precision="{"4" if prec=="single" else "8"}" Format="HDF">
        data.{time:04d}.{"flt" if prec=="single" else "dbl"}.h5:/Timestep_{time}/vars/VX3
       </DataItem>
     </Attribute>
   </Grid>
 </Domain>
</Xdmf>
'''
with open(f'data.{time:04d}.{"flt" if prec=="single" else "dbl"}.xmf', 'w') as txt:
    txt.write(xmf[1:])