{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm as cm\n",
    "import xarray as xr\n",
    "import dask\n",
    "import pandas as pd\n",
    "import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#observations directory\n",
    "obsdir = '/global/cscratch1/sd/avarble/eagles/observations/ena/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#constants\n",
    "G = 9.8      #gravity\n",
    "Cp = 1005.7  #specific heat of air at constant pressure\n",
    "Rd = 287.    #dry air constant\n",
    "Rv = 461.    #moist air constant\n",
    "lv = 2.477e6 #latent heat of vaporization at 10 C\n",
    "epsilon = Rd/Rv\n",
    "pres_const = 85000. #Pa (used by Bennartz (2007, JGR)), could use cloud top pressure, but shouldn't alter estimates much\n",
    "Q = 2.   #scattering efficiency\n",
    "k = 0.74 #drop size dispersion; Bennartz (2007, JGR) uses 0.8 +/- 0.1 but can be 0.5-0.9 depending on cloud type\n",
    "rho_liq = 1000. #liquid water density"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#loop through all 4-km pixel VISST satellite retrieval files to compute cloud drop concentration following Bennartz (2007, JGR)\n",
    "files = glob.glob(os.path.join(obsdir, 'visst/pix/enavisstpx2dm11minnisX1.c1.*.cdf'))\n",
    "files2 = sorted(files)\n",
    "for file in files2:\n",
    "    data = xr.open_dataset(file)\n",
    "    strdate = file.split('.')[2]\n",
    "    strtime = file.split('.')[3]\n",
    "    base_time = data['base_time']\n",
    "    time = data['time_offset']\n",
    "    lat = data['latitude']\n",
    "    lon = data['longitude']\n",
    "    phase = data['cloud_phase']               #cloud top phase\n",
    "    cod = data['cloud_visible_optical_depth'] #cloud optical depth\n",
    "    lwp = data['cloud_lwp_iwp']               #LWP or IWP depending on cloud top phase\n",
    "    ctt = data['cloud_top_temperature']       #cloud top temperature\n",
    "    \n",
    "    rho_air = pres_const/(Rd*ctt) #dry air density\n",
    "    es = 611.2*np.exp(17.62*(ctt-273.15)/(243.12 + ctt - 273.15))\n",
    "    ws = epsilon*es/(pres_const - es)\n",
    "    gamma_w = G*((1 + lv*ws/(Rd*ctt))/(Cp + lv**2*ws*epsilon/(Rd*ctt**2))) #moist adiabatic lapse rate in cloud\n",
    "    gamma_ad = (((epsilon + ws)*ws*lv*gamma_w)/(Rd*ctt**2) - (G*ws*pres_const/(Rd*ctt*(pres_const - es))))*rho_air #adiabatic LWC lapse rate in cloud\n",
    "    H = (2.*1e-3*lwp/(0.8*gamma_ad))**0.5 #cloud depth assuming 80% adiabatic fraction\n",
    "    Nd = 1e-6*(cod**3/k)*((2*(1e-3*lwp))**(-2.5))*((0.6*np.pi*Q)**(-3))*((3./(4.*np.pi*rho_liq))**(-2))*((0.8*gamma_ad)**0.5) #cloud droplet number (80% adiabatic)\n",
    "    H_ad = (2.*1e-3*lwp/(gamma_ad))**0.5 #cloud depth assuming 100% adiabatic\n",
    "    Nd_ad = 1e-6*(cod**3/k)*((2*(1e-3*lwp))**(-2.5))*((0.6*np.pi*Q)**(-3))*((3./(4.*np.pi*rho_liq))**(-2))*((gamma_ad)**0.5) #100% adiabatic cloud droplet number\n",
    "\n",
    "    H_array = np.array(H)\n",
    "    H_ad_array = np.array(H_ad)\n",
    "    Nd_array = np.array(Nd)\n",
    "    Nd_ad_array = np.array(Nd_ad)\n",
    "    \n",
    "    #set columns with IWP > 0 to NaN\n",
    "    ind = np.array(iwp > 0)\n",
    "    H_array[ind] = np.nan\n",
    "    H_ad_array[ind] = np.nan\n",
    "    Nd_array[ind] = np.nan\n",
    "    Nd_ad_array[ind] = np.nan\n",
    "    \n",
    "    #set columns with bad retrievals to NaN\n",
    "    ind = np.isinf(Nd_array)\n",
    "    H_array[ind] = np.nan\n",
    "    H_ad_array[ind] = np.nan\n",
    "    Nd_array[ind] = np.nan\n",
    "    Nd_ad_array[ind] = np.nan\n",
    "    \n",
    "    #write to file for analyses\n",
    "    ds = xr.Dataset({'base_time': ('time', np.arange(1)), 'base_time': base_time,\n",
    "                 'time_offset': ('time', np.arange(1)), 'time_offset': time,\n",
    "                 'CDNC': (('y','x'), np.float32(Nd_array)),\n",
    "                 'CDNC_adiabatic': (('y','x'), np.float32(Nd_ad_array)),\n",
    "                 'H': (('y','x'), np.float32(H_array)),\n",
    "                 'H_adiabatic': (('y','x'), np.float32(H_ad_array))})\n",
    "\n",
    "    ds['base_time'].attrs[\"ancillary_variables\"] = \"time_offset\"\n",
    "    ds['time_offset'].attrs[\"long_name\"] = \"Time offset from base_time\"\n",
    "    ds['time_offset'].attrs[\"ancillary_variables\"] = \"base_time\"\n",
    "    ds['longitude'].attrs[\"long_name\"] = \"Longitude\"\n",
    "    ds['longitude'].attrs[\"units\"] = \"degrees\"\n",
    "    ds['latitude'].attrs[\"long_name\"] = \"Latitude\"\n",
    "    ds['latitude'].attrs[\"units\"] = \"degrees\"\n",
    "    ds['CDNC'].attrs[\"long_name\"] = \"Cloud Droplet Concentration\"\n",
    "    ds['CDNC'].attrs[\"units\"] = \"cm-3\"\n",
    "    ds['CDNC'].attrs[\"description\"] = \"Retrieved following Bennartz using VISST product Meteosat data assuming adiabaticity = 0.8.\"\n",
    "    ds['CDNC_adiabatic'].attrs[\"long_name\"] = \"Adiabatic Cloud Droplet Concentration\"\n",
    "    ds['CDNC_adiabatic'].attrs[\"units\"] = \"cm-3\"\n",
    "    ds['CDNC_adiabatic'].attrs[\"description\"] = \"Retrieved following Bennartz using VISST product Meteosat data assuming adiabaticity = 1.\"\n",
    "    ds['H'].attrs[\"long_name\"] = \"Cloud Depth\"\n",
    "    ds['H'].attrs[\"units\"] = \"m\"\n",
    "    ds['H'].attrs[\"description\"] = \"Estimated using satellite-retrieved LWP and moist adiabatic lapse rate with adiabaticity = 0.8.\"\n",
    "    ds['H_adiabatic'].attrs[\"long_name\"] = \"Adiabatic Cloud Depth\"\n",
    "    ds['H_adiabatic'].attrs[\"units\"] = \"m\"\n",
    "    ds['H_adiabatic'].attrs[\"description\"] = \"Estimated using satellite-retrieved LWP and moist adiabatic lapse rate with adiabaticity = 1.\"\n",
    "\n",
    "    ds.attrs[\"description\"] = \"Cloud droplet concentration retrievals are only valid in overcast single layer liquid cloud situations with sufficiently high solar angles. The retrieval follows Bennartz with an adiabaticity of 0.8, except k is set to 0.74 rather than 0.8. The adiabatic retrievals assume adiabaticity = 1.\"\n",
    "    ds.attrs[\"contact\"] = \"Adam Varble, adam.varble@pnnl.gov\"\n",
    "    ds.attrs[\"date\"] = \"3 January 2022\"\n",
    "\n",
    "    outfile = 'enavisstpx2dm10minnisX1.cdnc.c1.'+strdate+'.'+strtime+'.cdf'\n",
    "    ds.to_netcdf(obsdir+'visst/cdnc/'+outfile, mode='w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MyEnv",
   "language": "python",
   "name": "myenv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
