{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Calculate and plot weighted means\n\nThis is an example of how to calculate a\nweighted average from the MET TBRG, ORG\nand PWD.  This also calculates the\naccumulated precipitation and displays it\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport xarray as xr\nfrom arm_test_data import DATASETS\n\nimport act\n\n# Specify dictionary of datastreams, variables, and weights\n# Note, all weights should add up to 1.\ncf_ds = {\n    'sgpmetE13.b1': {\n        'variable': [\n            'tbrg_precip_total',\n            'org_precip_rate_mean',\n            'pwd_precip_rate_mean_1min',\n        ],\n        'weight': [0.8, 0.15, 0.05],\n    }\n}\n\n# Other way to define cf_ds\n# cf_ds = {'sgpmetE13.b1': {'variable': ['tbrg_precip_total'], 'weight': [0.5]},\n#         'sgpmetE13.b1': {'variable': ['org_precip_rate_mean'], 'weight': [0.25]},\n#         'sgpmetE13.b1': {'variable': ['pwd_precip_rate_mean_1min'], 'weight': [0.25]}\n#        }\n\n# Get a list of filenames to use\nmet_wildcard_list = [\n    'sgpmetE13.b1.20190101.000000.cdf',\n    'sgpmetE13.b1.20190102.000000.cdf',\n    'sgpmetE13.b1.20190103.000000.cdf',\n    'sgpmetE13.b1.20190104.000000.cdf',\n    'sgpmetE13.b1.20190105.000000.cdf',\n    'sgpmetE13.b1.20190106.000000.cdf',\n    'sgpmetE13.b1.20190107.000000.cdf',\n]\n\nds = {}\nnew = {}\nout_units = 'mm/hr'\nfor d in cf_ds:\n    met_filenames = [DATASETS.fetch(file) for file in met_wildcard_list]\n    ds = act.io.arm.read_arm_netcdf(met_filenames)\n    # Loop through each variable and add to data list\n    new_da = []\n    for v in cf_ds[d]['variable']:\n        da = ds[v]\n        # Accumulate precip variables in new dataset\n        ds = act.utils.data_utils.accumulate_precip(ds, v)\n\n        # Convert units and add to dataarray list\n        units = da.attrs['units']\n        if units == 'mm':\n            da.attrs['units'] = 'mm/min'\n        da.values = act.utils.data_utils.convert_units(da.values, da.attrs['units'], out_units)\n        da = da.resample(time='1min').mean()\n        new_da.append(da)\n\n    # Depending on number of variables for each datastream, merge or create dataset\n    if len(new_da) > 1:\n        new_da = xr.merge(new_da)\n    else:\n        new_da = new_da[0].to_dataset()\n\n    # Add to dictionary for the weighting\n    cf_ds[d]['ds'] = new_da\n\n    # Add dataset to dictionary for plotting\n    new[d] = ds\n\n# Calculate weighted averages using the dict defined above\ndata = act.utils.data_utils.ts_weighted_average(cf_ds)\n\n# Add weighted mean to plotting object and calculate accumulation\nnew['weighted'] = data.to_dataset(name='weighted_mean')\nnew['weighted']['weighted_mean'].attrs['units'] = 'mm/hr'\nnew['weighted'] = act.utils.data_utils.accumulate_precip(new['weighted'], 'weighted_mean')\n\n# Plot the accumulations\ndisplay = act.plotting.TimeSeriesDisplay(new, figsize=(12, 8), subplot_shape=(1,))\ndisplay.plot('tbrg_precip_total_accumulated', dsname='sgpmetE13.b1', color='b', label='TBRG, 0.8')\ndisplay.plot(\n    'org_precip_rate_mean_accumulated',\n    dsname='sgpmetE13.b1',\n    color='g',\n    label='ORG 0.15',\n)\ndisplay.plot(\n    'pwd_precip_rate_mean_1min_accumulated',\n    dsname='sgpmetE13.b1',\n    color='y',\n    label='PWD 0.05',\n)\ndisplay.plot('weighted_mean_accumulated', dsname='weighted', color='k', label='Weighted Avg')\ndisplay.day_night_background('sgpmetE13.b1')\ndisplay.axes[0].legend()\nplt.show()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "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.14.3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}