# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## Overview

import numpy as np
import pandas as pd
import xarray as xr

# +
data = xr.DataArray(np.random.randn(2,3),
                    dims=('x', 'y'),
                    coords={'x': [10,20]})

data
# -

data.values

type(data.values)

data.dims

data.coords

data.attrs

# ### Indexing

# positional and by integer label, like numpy
data[0, :]

# loc or "location": positional and coordinate label, like pandas
data.loc[10]

# isel or "integer select": by dimension name and integer label
data.isel(x=0)

# sel or "select": by dimension name and coordinate label
data.sel(x=10)

# ### Attributes
# When setting up a DataArray, it's a good idea to set metadata attributes.

# +
data.attrs['long_name'] = 'random velocity'
data.attrs['units'] = 'meters/sec'
data.attrs['description'] = 'A random variable created as an example'
data.attrs['random_attribute'] = 123

data
# -

np.sin(data)

data.T

data.sum()

# ### GroupBy example

labels = xr.DataArray(['E', 'F', 'E'], [data.coords['y']], name='labels')
labels

data.groupby(labels).mean('y')

data.plot()

# ### Datasets

# +
ds = xr.Dataset({'foo': data,
                 'bar': ('x', [1, 2]),
                 'baz': np.pi})

ds
# -

# ### Saving and loading
# NetCDF is the recommended format.

ds.to_netcdf('example.nc')

x = xr.open_dataset('example.nc')
x

# ## DataArray

# ?xr.DataArray

# +
data = np.random.randn(4,3)
locs = ['IA', 'IL', 'IN']
times = pd.date_range('2000-01-01', periods=4)

da = xr.DataArray(data, coords=[times, locs], dims=['time', 'space'])

da
# -

xr.DataArray(data)

# It will return the coords
da['time']

da.coords['time']

# +
temp = xr.DataArray(data, dims=('time', 'space'))

temp['time']
# -

# ## Dataset

# +
latitude  = [f'lat_{x}' for x in range(5)]
longitude = [f'long_{x}' for x in range(7)]
time_stamps = [f'time_{x}' for x in range(10)]

temperature = np.random.randn(5, 7, 10) # lat, long, time
percipitation = np.random.randn(5, 7, 10) # lat, long, time

# +
# We can combine all the above data into a a DataArray where we would need to concatenate the arrays
da = xr.DataArray(data = np.stack((temperature, percipitation)),
             coords = {
                 'type': ['temperature', 'percipitation'],
                 'latitude': latitude,
                 'longitude': longitude,
                 'time': time_stamps,
             },
             dims = ('type', 'latitude', 'longitude', 'time'))

da
# -

# In the above example we kind of introduced complexity as we had to stack two arrays and introduce a new dimension named `type`. We can see that in the above example for both the arrays (temperature and percipitation) they share the same dims. `Dataset` can help in this case.

# +
# First let's create temperature and percipitation DataArrays
temp_da = xr.DataArray(temperature,
                       coords = {
                           'latitude': latitude,
                           'longitude': longitude,
                           'time': time_stamps,
                       },
                       dims = ('latitude', 'longitude', 'time'))

percip_da = xr.DataArray(percipitation,
                       coords = {
                           'latitude': latitude,
                           'longitude': longitude,
                           'time': time_stamps,
                       },
                       dims = ('latitude', 'longitude', 'time'))

# Create Dataset
ds = xr.Dataset({
        'temperature': temp_da,
        'percipitation': percip_da,
     })

ds
# -

ds['temperature']

# ## Indexing and Selecting Data

# +
latitude  = [f'lat_{x}' for x in range(2)]
longitude = [f'long_{x}' for x in range(3)]

temperature = np.random.randn(2, 3) # lat, long, time

da = xr.DataArray(data = temperature,
                 coords = {
                     'latitude': latitude,
                     'longitude': longitude,
                 },
                 dims = ('latitude', 'longitude'))

da
# -

da.values

da[0,2].values

da.loc['lat_0', 'long_1'].values

da.isel(latitude=0).values

# The av=bove is same as 
da[dict(latitude=0)].values

da.sel(latitude='lat_1').values

# The above is same as
da.loc[dict(latitude='lat_1')].values

# ### Indexing with dimension names

# +
latitude  = [f'lat_{x}' for x in range(2)]
longitude = [f'long_{x}' for x in range(3)]
time_stamps = [f'time_{x}' for x in range(5)]

temperature = np.random.randn(2, 3, 5) # lat, long, time

da = xr.DataArray(data = temperature,
                 coords = {
                     'latitude': latitude,
                     'longitude': longitude,
                     'time': time_stamps,
                 },
                 dims = ('latitude', 'longitude', 'time'))


da.values
# -

da.isel(latitude=0, time=slice(None, 2)).values

da.sel(latitude='lat_0', time=slice(None, 'time_1')).values

# ### Nearest neighbor lookups

da = xr.DataArray([1, 2, 3], [('x', [0, 1, 2])])
da

da.sel(x=[1.1, 1.9], method='nearest').values

da.sel(x=0.1, method='backfill').values

da.sel(x=[0.5, 1, 1.9, 2, 2.5, 1.3], method='pad').values

da.reindex(x=[1.1, 1.5], method='nearest', tolerance=0.2).values



StarcoderdataJupyterScriptsDedupFiltered