third_order_conversion.py

from numba import njit, prange, set_num_threads
import netCDF4 as nc
import numpy as np
from glob import glob
from math import erf
import argparse
import os

# Note: check division by zero errors.

projfac = 4  # Subgrid zoom factor
r_limit_fac = 3.0
#r_limit_fac = 1.0  # Alternative: individual ellipsoid visualisation
nthreads = 12
set_num_threads(nthreads)

parser = argparse.ArgumentParser()
parser.add_argument("input_file_name", type=str)
args = parser.parse_args()
input_file_name = args.input_file_name
file_root, file_ext = os.path.splitext(input_file_name)

if not (file_ext == ".nc"):
    raise argparse.ArgumentTypeError('Argument filename must end on ".nc"')

ds_nc = nc.Dataset(input_file_name)

# set up spatial arrays
extent = ds_nc.extent
origin = ds_nc.origin
ncells = ds_nc.ncells
len_condense = ds_nc.groups['physical_quantities'].scale_height
q_scale = ds_nc.groups['physical_quantities'].saturation_specific_humidity_at_ground_level

# initialisation
nx = ncells[0]
ny = ncells[1]
nz = ncells[2]  # Don't add 1 (this is non-periodic) here
# projection grid
nxproj = nx * projfac
nyproj = ny * projfac
nzproj = nz * projfac + 1  # But do add 1 here
dx = extent[0] / nx
dy = extent[1] / ny
dz = extent[2] / nz
dx_project = dx / projfac
dy_project = dy / projfac
dz_project = dz / projfac
dxi_project = projfac / dx  # Inverse of projection grid distance
dyi_project = projfac / dy
dzi_project = projfac / dz

# Projection grid coordinates
xp = np.linspace(origin[0], origin[0] + extent[0], nxproj + 1)[0:nxproj]
yp = np.linspace(origin[1], origin[1] + extent[1], nyproj + 1)[0:nyproj]
zp = np.linspace(origin[2], origin[2] + extent[2], nzproj)

ipi = 1.0 / np.pi
# Summation fields
xyz_field = np.zeros((nxproj, nyproj, nzproj), np.double)
xyz_volg = np.zeros((nxproj, nyproj, nzproj), np.double)  # Normalisation factor

# This function operates on small subdomains
# Does not yet work in parallel
@njit(parallel=True)
def add_data(
    h,
    x,
    y,
    z,
    bb1,
    bb2,
    bb3,
    bb4,
    bb5,
    vol,
    xyz_field_thread,
    xyz_volg_thread,
    xlowerfile,
    ylowerfile,
):
    # Construction below implements something that is a lot like OpenMP reduction in numba
    # Each thread adds to its own field, and only later we add threads together in add_fields
    for i_thread in prange(nthreads):
        bb = np.zeros((3, 3), np.double)
        bbinv = np.zeros((3, 3), np.double)
        xarr = np.zeros((3, 1), np.double)
        evals = np.zeros(3, np.double)
        for i in range(len(h)):
            if (
                not i % nthreads == i_thread
            ):  # Use modulus here, not sure if contiguous blocks could be better
                continue
            vol_now = vol[i]
            r_now = (vol_now * 0.75 * ipi) ** (1.0 / 3.0)
            # Maximum radius for parcel projection, and related constants
            x_now = x[i]
            y_now = y[i]
            z_now = z[i]
            scalar_now = h[i]
            b1 = bb1[i]
            b2 = bb2[i]
            b3 = bb3[i]
            b4 = bb4[i]
            b5 = bb5[i]
            abc = 0.75 * vol_now * ipi
            b6 = (
                abc * abc - b3 * (b2 * b5 - b3 * b4) + b1 * b5 * b5 - b2 * b3 * b5
            ) / (b1 * b4 - b2 * b2)
            bb[0, 0] = b1
            bb[0, 1] = b2
            bb[0, 2] = b3
            bb[1, 0] = b2
            bb[1, 1] = b4
            bb[1, 2] = b5
            bb[2, 0] = b3
            bb[2, 1] = b5
            bb[2, 2] = b6
            evals[:] = np.linalg.eigvalsh(bb)
            bbinv[:, :] = np.linalg.inv(bb)
            anisotropy_fact = np.sqrt(np.nanmax(evals)) / (
                (evals[0] * evals[1] * evals[2]) ** (1.0 / 6.0)
            )
            r_max_safety = r_limit_fac * anisotropy_fact
            r_max = r_max_safety * r_now
            # Determine which part of the grid the parcel contributes to
            xlower = int(dxi_project * ((x_now - r_max) - origin[0])) + 1
            xupper = int(dxi_project * ((x_now + r_max) - origin[0]))
            ylower = int(dyi_project * ((y_now - r_max) - origin[1])) + 1
            yupper = int(dyi_project * ((y_now + r_max) - origin[1]))
            zlower = max(int(dzi_project * ((z_now - r_max) - origin[2])), 0)
            zupper = min(int(dzi_project * ((z_now + r_max) - origin[2])), nzproj - 1)
            norm_fact = 1.0
            # Include upper values in loops (converted from Fortran)
            for ii in range(xlower, xupper + 1):
                for jj in range(ylower, yupper + 1):
                    for kk in range(zlower, zupper + 1):
                        xdist = ii * dx_project + origin[0] - x_now
                        ydist = jj * dy_project + origin[1] - y_now
                        zdist = kk * dz_project + origin[2] - z_now
                        dist2 = xdist * xdist + ydist * ydist + zdist * zdist
                        if dist2 < r_max * r_max:
                            xarr[0, 0] = xdist
                            xarr[1, 0] = ydist
                            xarr[2, 0] = zdist
                            r_act = np.sqrt(
                                np.dot(np.dot(xarr.transpose(), bbinv), xarr,)[0, 0]
                            )
                            # r_act = np.sqrt(dist2)/r_now
                            if r_act < r_limit_fac:
                                # Correct y and x location with respect to file-based grids
                                # Do not correct for cyclic boundary conditions yet, this is done in add_fields
                                ii_corr = ii - xlowerfile
                                jj_corr = jj - ylowerfile
                                r_act3 = r_act * r_act * r_act
                                #r_fact = (
                                #    r_act < 1.0
                                #)  # Alternative: individual ellipsoid visulatisation.
                                r_fact = np.exp(-r_act3)
                                xyz_field_thread[i_thread, ii_corr, jj_corr, kk] += (
                                    norm_fact * r_fact * scalar_now
                                )
                                xyz_volg_thread[i_thread, ii_corr, jj_corr, kk] += (
                                    norm_fact * r_fact
                                )


# Sum up thread-specific sums on local coordinates to the correct global coordinates
@njit()
def add_fields(
    xyz_field,
    xyz_volg,
    xyz_field_thread_file,
    xyz_volg_thread_file,
    xlowerfile,
    ylowerfile,
    nx_file,
    ny_file,
):
    for i_thread in range(nthreads):
        # Parallelise this routine in x, rather than over threads, so that conflicts are avoided
        for ii in range(nx_file):
            for jj in range(ny_file):
                for kk in range(nzproj):
                    itarg = (ii + xlowerfile) % nxproj  # Periodic BCs
                    jtarg = (jj + ylowerfile) % nyproj  # Periodic BCs
                    xyz_field[itarg, jtarg, kk] = (
                        xyz_field[itarg, jtarg, kk]
                        + xyz_field_thread_file[i_thread, ii, jj, kk]
                    )
                    xyz_volg[itarg, jtarg, kk] = (
                        xyz_volg[itarg, jtarg, kk]
                        + xyz_volg_thread_file[i_thread, ii, jj, kk]
                    )


h = ds_nc.variables["humidity"][0,:]
x = ds_nc.variables["x_position"][0,:]
y = ds_nc.variables["y_position"][0,:]
z = ds_nc.variables["z_position"][0,:]
vol = ds_nc.variables["volume"][0,:]
bb1 = ds_nc.variables["B11"][0,:]
bb2 = ds_nc.variables["B12"][0,:]
bb3 = ds_nc.variables["B13"][0,:]
bb4 = ds_nc.variables["B22"][0,:]
bb5 = ds_nc.variables["B23"][0,:]

# Local domain min and max values
xminval = np.nanmin(x)
xmaxval = np.nanmax(x)
yminval = np.nanmin(y)
ymaxval = np.nanmax(y)
# Local domain projection coordinates boundaries
# Add one (round up rather than down), as values are zero at r_max distance
xlowerfile = int(dxi_project * ((xminval - 3.0 * dx) - origin[0])) + 1
xupperfile = int(dxi_project * ((xmaxval + 3.0 * dx) - origin[0]))
ylowerfile = int(dyi_project * ((yminval - 3.0 * dy) - origin[1])) + 1
yupperfile = int(dyi_project * ((ymaxval + 3.0 * dy) - origin[1]))
# Size of local projection domain, which includes both xupper and xlower etc.
nx_file = xupperfile - xlowerfile + 1
ny_file = yupperfile - ylowerfile + 1
# Initialise "thread-based" fields for this file
xyz_field_thread_file = np.zeros((nthreads, nx_file, ny_file, nzproj), np.double)
xyz_volg_thread_file = np.zeros((nthreads, nx_file, ny_file, nzproj), np.double)
add_data(
    h,
    x,
    y,
    z,
    bb1,
    bb2,
    bb3,
    bb4,
    bb5,
    vol,
    xyz_field_thread_file,
    xyz_volg_thread_file,
    xlowerfile,
    ylowerfile,
)
add_fields(
    xyz_field,
    xyz_volg,
    xyz_field_thread_file,
    xyz_volg_thread_file,
    xlowerfile,
    ylowerfile,
    nx_file,
    ny_file,
)

with open("tracker.txt", "a") as f:
    print("added fields", file=f)

del xyz_field_thread_file, xyz_volg_thread_file

# Final calculation of liquid water field
# Normalize h field, and convert to LWC
# See MONC/MPIC intercomparison appendix B
# This is like strategy C
@njit()
def final_calc(field, volg, zp):
    for ii in range(np.shape(field)[0]):
        for jj in range(np.shape(field)[1]):
            for kk in range(np.shape(field)[2]):
                if volg[ii, jj, kk] > 0.0:
                    field_temp = field[ii, jj, kk] / volg[ii, jj, kk]
                    volg_temp = max(
                        field_temp
                        - q_scale * np.exp((origin[2] - zp[kk]) / len_condense),
                        0.0,
                    )
                else:
                    field_temp = np.nan
                    volg_temp = np.nan
                field[ii, jj, kk] = field_temp
                volg[ii, jj, kk] = volg_temp


# store total water in field
# store liquid water in volg
final_calc(xyz_field, xyz_volg, zp)

with open("tracker.txt", "a") as f:
    print("div ready", file=f)

# Write to netcdf file
def write_field(field, field_name, out_file, mode):
    ncfile = nc.Dataset(out_file, mode, format="NETCDF4")
    if mode == "w":
        ncfile.createDimension("X", nxproj)
        ncfile.createDimension("Y", nyproj)
        ncfile.createDimension("Z", nzproj)
        nc_xp = ncfile.createVariable("X", "f8", ("X",))
        nc_yp = ncfile.createVariable("Y", "f8", ("Y",))
        nc_zp = ncfile.createVariable("Z", "f8", ("Z",))
        nc_xp.units = "m"
        nc_yp.units = "m"
        nc_zp.units = "m"
        nc_xp.axis = 'X' # Optional
        nc_xp.standard_name = 'projection_x_coordinate'
        nc_xp.long_name = 'x-coordinate in projected coordinate system'
        nc_yp.axis = 'Y' # Optional
        nc_yp.standard_name = 'projection_y_coordinate'
        nc_yp.long_name = 'y-coordinate in projected coordinate system'
        nc_zp.axis = 'Z' # Optional
        nc_zp.standard_name = 'height'
        nc_zp.long_name = 'height coordinate in projected coordinate system'
        nc_xp[:] = xp
        nc_yp[:] = yp
        nc_zp[:] = zp
    nc_hl = ncfile.createVariable(field_name, np.dtype("f4").char, ("Z", "Y", "X"))
    nc_hl.units = "1"
    nc_hl[:, :, :] = np.transpose(field[:, :, :])
    nc_hl.long_name = field_name
    ncfile.close()


write_field(xyz_field, "hh", file_root + "_g3.nc", "w")
write_field(xyz_volg, "hl", file_root + "_g3.nc", "a")