title: python get month of maximum value xarray
body: <h1>How to get the month of maximum runoff</h1>
  <p>I want to get the month of maximum runoff for each year, and for the time series as a whole. The idea is to characterise global seasonality by looking at the month of max runoff. I then want to try and consider whether each pixel has a unimodal or bimodal regime.</p>
  <p>I want to create a map like the one in the Pangeo Examples <a href="https://github.com/pangeo-data/pangeo_cmip6_examples" rel="nofollow noreferrer">here</a>.</p>
  <p><a href="https://i.stack.imgur.com/1w80i.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/1w80i.png" alt="Example image" /></a></p>
  <p>What this shows is the hour of maximum precipitation. I want to show the MONTH of maximum runoff (as an integer).</p>
  <h2>Getting the data</h2>
  <p>Here I download the <a href="https://www.research-collection.ethz.ch/handle/20.500.11850/324386" rel="nofollow noreferrer">GRUN runoff data</a> and create an xarray object.
  <strong>NOTE:</strong> The dataset here is &gt;1GB. I am using it to make this example entirely reproducible.</p>
  <pre class="lang-py prettyprint-override"><code># get the data
  import subprocess
  command = &quot;&quot;&quot;
  wget -O grun.nc https://www.research-collection.ethz.ch/bitstream/handle/20.500.11850/324386/GRUN_v1_GSWP3_WGS84_05_1902_2014.nc?sequence=1&amp;isAllowed=y
  &quot;&quot;&quot;
  import os
  if not os.path.exists('grun.nc'):
    process = subprocess.Popen(command.split(), stdout=subprocess.PIPE)
    output, error = process.communicate()
  
  # read the data
  import xarray as xr
  ds = xr.open_dataset('grun.nc')
  
  # select a subset so we can work with it more quickly
  ds = ds.isel(time=slice(-100,-1))
  ds
  
  Out[]:
  &lt;xarray.Dataset&gt;
  Dimensions:  (lat: 360, lon: 720, time: 99)
  Coordinates:
    * lon      (lon) float64 -179.8 -179.2 -178.8 -178.2 ... 178.8 179.2 179.8
    * lat      (lat) float64 -89.75 -89.25 -88.75 -88.25 ... 88.75 89.25 89.75
    * time     (time) datetime64[ns] 2006-09-01 2006-10-01 ... 2014-11-01
  Data variables:
      Runoff   (time, lat, lon) float32 ...
  Attributes:
      title:                   GRUN
      version:                 GRUN 1.0
      meteorological_forcing:  GSWP3
      temporal_resolution:     monthly
      spatial_resolution:      0.5x0.5
      crs:                     WGS84
      proj4:                   +proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs
      EPSG:                    4326
      references:              Ghiggi et al.,2019. GRUN: An observation-based g...
      authors:                 Gionata Ghiggi; Lukas Gudmundsson
      contacts:                gionata.ghiggi@gmail.com; lukas.gudmundsson@env....
      institution:             Land-Climate Dynamics, Institute for Atmospheric...
      institution_id:          IAC ETHZ
  
  </code></pre>
  <h2>What I have tried</h2>
  <p>I have nan values so I can't just apply an <code>argmax()</code> to the dataset. I use the same approach as @jhamman <a href="https://stackoverflow.com/a/50499427/9940782">here</a> combined with the Pangeo Examples above. I'm not entirely sure what this is giving me but it seems to be giving me</p>
  <pre class="lang-py prettyprint-override"><code># Apply argmax where you have NAN values
  def my_func(ds, dim=None):
      return ds.isel(**{dim: ds['Runoff'].argmax(dim)})
  
  mask = ds['Runoff'].isel(time=0).notnull()  # determine where you have valid data
  ds2 = ds.fillna(-9999)  # fill nans with a missing flag of some kind
  new = ds2.reset_coords(drop=True).groupby('time.month').apply(my_func, dim='time').where(mask)  # do the groupby operation/reduction and reapply the mask
  new
  
  Out[]:
  &lt;xarray.Dataset&gt;
  Dimensions:  (lat: 360, lon: 720, month: 12)
  Coordinates:
    * lon      (lon) float64 -179.8 -179.2 -178.8 -178.2 ... 178.8 179.2 179.8
    * lat      (lat) float64 -89.75 -89.25 -88.75 -88.25 ... 88.75 89.25 89.75
    * month    (month) int64 1 2 3 4 5 6 7 8 9 10 11 12
  Data variables:
      Runoff   (month, lat, lon) float32 nan nan nan nan nan ... nan nan nan nan
  Attributes:
      title:                   GRUN
      version:                 GRUN 1.0
      meteorological_forcing:  GSWP3
      temporal_resolution:     monthly
      spatial_resolution:      0.5x0.5
      crs:                     WGS84
      proj4:                   +proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs
      EPSG:                    4326
      references:              Ghiggi et al.,2019. GRUN: An observation-based g...
      authors:                 Gionata Ghiggi; Lukas Gudmundsson
      contacts:                gionata.ghiggi@gmail.com; lukas.gudmundsson@env....
      institution:             Land-Climate Dynamics, Institute for Atmospheric...
      institution_id:          IAC ETHZ
  
  </code></pre>
  <p>This gives me</p>
  <pre class="lang-py prettyprint-override"><code>import matplotlib.pyplot as plt
  fig,ax = plt.subplots(figsize=(12,8))
  new.Runoff.sel(month=10).plot(ax=ax,  cmap='twilight')
  </code></pre>
  <p><a href="https://i.stack.imgur.com/X3mlE.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/X3mlE.png" alt="My output currently" /></a></p>
  <h2> Ideal Output</h2>
  <h3>What I want is for the value of each Pixel to be the month of maximum Runoff.</h3>
  <p>Happy to convert to <code>pandas</code> if necessary.</p>
  <p>So I would end up with a xr.Dataset with the integer for the month of maximum runoff. Ideally, it would be great to also have the month of maximum runoff over time so I can also see the way that this seasonality has changed.</p>
  <pre class="lang-py prettyprint-override"><code>&lt;xarray.Dataset&gt;
  Dimensions:  (lat: 360, lon: 720)
  Coordinates:
    * lon      (lon) float64 -179.8 -179.2 -178.8 -178.2 ... 178.8 179.2 179.8
    * lat      (lat) float64 -89.75 -89.25 -88.75 -88.25 ... 88.75 89.25 89.75
  Data variables:
      Month_of_max (lat, lon) int32 ...
      
  # OR EVEN BETTER
  &lt;xarray.Dataset&gt;
  Dimensions:  (lat: 360, lon: 720, Year: 10)
  Coordinates:
    * lon      (lon) float64 -179.8 -179.2 -178.8 -178.2 ... 178.8 179.2 179.8
    * lat      (lat) float64 -89.75 -89.25 -88.75 -88.25 ... 88.75 89.25 89.75
    * year     (year) float64 2010 2011 2012 2013 ... 
  Data variables:
      Month_of_max (lat, lon, year) int32 ...
  </code></pre>