Rachel-20240116.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2361412d-d7fe-45c6-923a-0c626f6ea541",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import modules\n",
    "import six\n",
    "import mule\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import os\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f29ae6e1-0e06-43ba-811b-ef4ba4784d5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set some config\n",
    "restart_filepath = '/g/data/p66/rml599/LandUse-tests/pre-industrial-essential.astart'\n",
    "stash_master_filepath = '/g/data/access/projects/access/umdir/vn10.9/ctldata/STASHmaster/STASHmaster_A'\n",
    "vegfrac_filepath = '/g/data/p66/rml599/LandUse-tests/vegfrac-dom-tree.nc' # step 2\n",
    "stash_range = range(851, 860+1)\n",
    "step2 = True\n",
    "levels = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9d903a35-6b7b-48ed-904d-3fa96f677c8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/g/data/hh5/public/apps/miniconda3/envs/analysis3-23.10/lib/python3.10/site-packages/mule/stashmaster.py:259: UserWarning: \n",
      "Unable to load STASHmaster from version string, path does not exist\n",
      "Path: $UMDIR/vn7.3/ctldata/STASHmaster/STASHmaster_A\n",
      "Please check that the value of mule.stashmaster.STASHMASTER_PATH_PATTERN is correct for your site/configuration\n",
      "  warnings.warn(msg)\n",
      "/g/data/hh5/public/apps/miniconda3/envs/analysis3-23.10/lib/python3.10/site-packages/mule/validators.py:198: UserWarning: \n",
      "File: /g/data/p66/rml599/LandUse-tests/pre-industrial-essential.astart\n",
      "Incorrect dataset_type (found 1, should be one of (3,))\n",
      "  warnings.warn(msg)\n"
     ]
    }
   ],
   "source": [
    "# Load the file\n",
    "restart = mule.FieldsFile.from_file(restart_filepath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4f35aac6-79c7-4c9d-897f-0aec2e180173",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_field(restart, stash_code, num_levels):\n",
    "\n",
    "    fields = list()\n",
    "    \n",
    "    for field in restart.fields:\n",
    "        if field.lbuser4 == stash_code and field.lbuser5 in range(1, num_levels+1):\n",
    "            fields.append(field)\n",
    "\n",
    "    return fields\n",
    "\n",
    "def to_numpy(field_list):\n",
    "    data = None\n",
    "    num_levels = len(field_list)\n",
    "\n",
    "    for field in field_list:\n",
    "\n",
    "        _data = field.get_data()\n",
    "\n",
    "        if data is None:\n",
    "            shape = (num_levels, *_data.shape)\n",
    "            data = np.zeros(shape)\n",
    "\n",
    "        ix = field.lbuser5 - 1\n",
    "        data[ix, :, :] = _data\n",
    "\n",
    "    return data\n",
    "\n",
    "def to_dataarray(field_list):\n",
    "\n",
    "    data = None\n",
    "    num_levels = len(field_list)\n",
    "\n",
    "    for field in field_list:\n",
    "\n",
    "        _data = field.get_data()\n",
    "\n",
    "        if data is None:\n",
    "            shape = (num_levels, *_data.shape)\n",
    "            data = np.zeros(shape)\n",
    "\n",
    "        ix = field.lbuser5 - 1\n",
    "        data[ix, :, :] = _data\n",
    "\n",
    "    da = xr.DataArray(\n",
    "        data,\n",
    "        dims=('z','y','x')\n",
    "    )\n",
    "    \n",
    "    return da\n",
    "    \n",
    "# Extract data from UM file.\n",
    "def extract_data_array(restart, stash_code, num_levels=5, do_mean=True):\n",
    "\n",
    "    data = None\n",
    "    \n",
    "    for field in restart.fields:\n",
    "        if field.lbuser4 == stash_code and field.lbuser5 in range(1, num_levels+1):\n",
    "\n",
    "            # First encounter, create the data object\n",
    "            if data is None:\n",
    "                shape = (num_levels, *field.get_data().shape)\n",
    "                data = np.zeros(shape)\n",
    "\n",
    "            ix = field.lbuser5 - 1\n",
    "            data[ix, :, :] = field.get_data()\n",
    "\n",
    "    # Create a data array\n",
    "    da = xr.DataArray(\n",
    "        data,\n",
    "        dims=('z', 'y', 'x')\n",
    "    )\n",
    "\n",
    "    if do_mean is False:\n",
    "        return da\n",
    "\n",
    "    # Nan the fill value\n",
    "    # fill_value = da.min()\n",
    "    # Leaving this out makes ocean super negative and the data invisible (only use for plotting!!)\n",
    "    # da = da.where(da != fill_value)\n",
    "    \n",
    "    da_mean = da.mean('z')\n",
    "\n",
    "    return da_mean\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1451c04c-7693-4513-8384-da3ab2d53ba4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the vegetation fracion\n",
    "vf = extract_field(restart, 216, levels)\n",
    "\n",
    "# Get the missing value\n",
    "# https://github.com/metomi/mule/blob/master/mule/docs/source/user_guide_basics.rst\n",
    "fill_value = vf[0].bmdi\n",
    "\n",
    "# # Convert to data array\n",
    "vf = to_dataarray(vf)\n",
    "vf_provided = xr.open_dataset(vegfrac_filepath)\n",
    "\n",
    "vf_provided = vf_provided.squeeze(drop=True).rename(vegtype='z', longitude='x', latitude='y').fraction.isel(z=range(0, levels))\n",
    "\n",
    "vfsum = vf.sum('z')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "13b72cf2-3518-4383-8338-1c785ddce00d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "851\n",
      "852\n",
      "853\n",
      "854\n",
      "855\n",
      "856\n",
      "857\n",
      "858\n",
      "859\n",
      "860\n"
     ]
    }
   ],
   "source": [
    "providers = dict()\n",
    "\n",
    "# Extract/modify the stashcodes for 851-860\n",
    "for stash_code in stash_range:\n",
    "    \n",
    "    print(stash_code)\n",
    "\n",
    "    # Extract the data\n",
    "    data = extract_field(restart, stash_code, num_levels=levels)\n",
    "\n",
    "    # Convert it to data array\n",
    "    data = to_dataarray(data)\n",
    "    \n",
    "    # Multiply the vegetation fraction\n",
    "    data_w = data * vf\n",
    "\n",
    "    # Sum across the dimension\n",
    "    data_ws = data_w.sum('z')\n",
    "\n",
    "    # Divide by the sum of the vegfrac, if nonzero\n",
    "    data_wsd = xr.where(vfsum == 0.0, 0.0, data_ws / vfsum)\n",
    "\n",
    "    # Re-establish the fill value\n",
    "    data_wsd = data_wsd.where(data_wsd != fill_value).fillna(fill_value)\n",
    "\n",
    "    # Create a data provider for this new field\n",
    "    providers[stash_code] = mule.ArrayDataProvider(data_wsd.data[:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9b2be5dd-9d7d-4f39-a715-71c3746749d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updating stash 216 level 1\n",
      "Updating stash 216 level 2\n",
      "Updating stash 216 level 3\n",
      "Updating stash 216 level 4\n",
      "Updating stash 216 level 5\n",
      "Updating stash 835 level 1\n",
      "Updating stash 835 level 2\n",
      "Updating stash 835 level 3\n",
      "Updating stash 835 level 4\n",
      "Updating stash 835 level 5\n",
      "Updating stash 851 level 1\n",
      "Updating stash 851 level 2\n",
      "Updating stash 851 level 3\n",
      "Updating stash 851 level 4\n",
      "Updating stash 851 level 5\n",
      "Updating stash 852 level 1\n",
      "Updating stash 852 level 2\n",
      "Updating stash 852 level 3\n",
      "Updating stash 852 level 4\n",
      "Updating stash 852 level 5\n",
      "Updating stash 853 level 1\n",
      "Updating stash 853 level 2\n",
      "Updating stash 853 level 3\n",
      "Updating stash 853 level 4\n",
      "Updating stash 853 level 5\n",
      "Updating stash 854 level 1\n",
      "Updating stash 854 level 2\n",
      "Updating stash 854 level 3\n",
      "Updating stash 854 level 4\n",
      "Updating stash 854 level 5\n",
      "Updating stash 855 level 1\n",
      "Updating stash 855 level 2\n",
      "Updating stash 855 level 3\n",
      "Updating stash 855 level 4\n",
      "Updating stash 855 level 5\n",
      "Updating stash 856 level 1\n",
      "Updating stash 856 level 2\n",
      "Updating stash 856 level 3\n",
      "Updating stash 856 level 4\n",
      "Updating stash 856 level 5\n",
      "Updating stash 857 level 1\n",
      "Updating stash 857 level 2\n",
      "Updating stash 857 level 3\n",
      "Updating stash 857 level 4\n",
      "Updating stash 857 level 5\n",
      "Updating stash 858 level 1\n",
      "Updating stash 858 level 2\n",
      "Updating stash 858 level 3\n",
      "Updating stash 858 level 4\n",
      "Updating stash 858 level 5\n",
      "Updating stash 859 level 1\n",
      "Updating stash 859 level 2\n",
      "Updating stash 859 level 3\n",
      "Updating stash 859 level 4\n",
      "Updating stash 859 level 5\n",
      "Updating stash 860 level 1\n",
      "Updating stash 860 level 2\n",
      "Updating stash 860 level 3\n",
      "Updating stash 860 level 4\n",
      "Updating stash 860 level 5\n"
     ]
    }
   ],
   "source": [
    "# Modify the file\n",
    "new_restart = restart.copy()\n",
    "\n",
    "for field in restart.fields:\n",
    "\n",
    "    stash_code = field.lbuser4\n",
    "    level = field.lbuser5\n",
    "\n",
    "    # If it is one of the fields we want to modify, do it\n",
    "    if stash_code in providers.keys() and level <= levels:\n",
    "        # new_restart.fields.append(operators[stash_code])\n",
    "        print(f'Updating stash {stash_code} level {level}')\n",
    "        field.set_data_provider(providers[stash_code])\n",
    "    \n",
    "    if stash_code in [216, 835] and level <= 5:\n",
    "        print(f'Updating stash {stash_code} level {level}')\n",
    "        provider = mule.ArrayDataProvider(vf_provided.data[level-1,:,:])\n",
    "        field.set_data_provider(provider)\n",
    "    \n",
    "    new_restart.fields.append(field)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "0d554810-a8b6-4036-ac80-4d02e6f5a3e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Intercept the write function to disable validation\n",
    "def to_file(self, output_file_or_path):\n",
    "        \"\"\"\n",
    "        Write to an output file or path.\n",
    "\n",
    "        Args:\n",
    "            * output_file_or_path (string or file-like):\n",
    "                An open file or filepath. If a path, it is opened and\n",
    "                closed again afterwards.\n",
    "\n",
    "        .. Note::\n",
    "            As part of this the \"validate\" method will be called. For the\n",
    "            base :class:`UMFile` class this does nothing, but sub-classes\n",
    "            may override it to provide specific validation checks.\n",
    "\n",
    "        \"\"\"\n",
    "        # Call validate - to ensure the file about to be written out doesn't\n",
    "        # contain obvious errors.  This is done here before any new file is\n",
    "        # created so that we don't create a blank file if the validation fails\n",
    "        if isinstance(output_file_or_path, six.string_types):\n",
    "            self.validate(filename=output_file_or_path, warn=True)\n",
    "        else:\n",
    "            self.validate(filename=output_file_or_path.name, warn=True)\n",
    "\n",
    "        if isinstance(output_file_or_path, six.string_types):\n",
    "            with open(output_file_or_path, 'wb') as output_file:\n",
    "                self._write_to_file(output_file)\n",
    "        else:\n",
    "            self._write_to_file(output_file_or_path)\n",
    "\n",
    "new_restart.to_file = to_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "9c449a69-35a3-406c-be47-d8c8b7af37c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/g/data/hh5/public/apps/miniconda3/envs/analysis3-23.10/lib/python3.10/site-packages/mule/validators.py:198: UserWarning: \n",
      "File: /g/data/p66/bjs581/mean5_VFW_851-860_20240204_101944.astart\n",
      "Incorrect dataset_type (found 1, should be one of (3,))\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/g/data/p66/bjs581/mean5_VFW_851-860_20240204_101944.astart\n"
     ]
    }
   ],
   "source": [
    "# Write out the data.\n",
    "stash_start, stash_end = stash_range[0], stash_range[-1]\n",
    "timestamp = datetime.utcnow().strftime('%Y%m%d_%H%M%S')\n",
    "\n",
    "output_filename = f'mean5_VFW_{stash_start}-{stash_end}_{timestamp}.astart'\n",
    "\n",
    "# if step2:\n",
    "#     output_filename = f'mean5_VFWDOM_{stash_start}-{stash_end}_{timestamp}.astart'    \n",
    "\n",
    "output_dir = '/g/data/p66/bjs581'\n",
    "output_filepath = f'{output_dir}/{output_filename}'\n",
    "\n",
    "new_restart.to_file(new_restart, output_file_or_path=output_filepath)\n",
    "\n",
    "print(output_filepath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b75cc915-9236-4b47-a19a-ac7224606ab7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:analysis3-23.10]",
   "language": "python",
   "name": "conda-env-analysis3-23.10-py"
  },
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}