# Multi-dimensional Coordinates¶

In [ ]:
import pathlib
from collections import defaultdict

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

import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm,Normalize
import matplotlib.cm as cm

from bluesky_tutorial_utils import nexus, fetch

In [ ]:
# Un-comment this for better-looking plots on high-resolution screens.
# %config InlineBackend.figure_format = 'retina'

In [ ]:
fetch.rsoxs_simulation_data();

In [ ]:
fname = './rsoxs_simulation_data/512-512-128-5.0-40.0-00285-0360.nxs'
da_img


## Extra: Multidimensional Coordinates¶

First let's look at our initial 2D array

In [ ]:
da_img


First we make a $Q=\sqrt{Q_{x}^2 + Q_{y}^2}$ array

In [ ]:
Q = np.sqrt(da_img.Qx*da_img.Qx + da_img.Qy*da_img.Qy)
Q.name = 'Q'
Q.plot()
Q


Then we make a $\chi= \tan^{-1}\left(\frac{Q_{y}}{Q_{x}}\right)$ array

In [ ]:
CHI = np.arctan2(da_img.Qy,da_img.Qx)*180/np.pi
CHI.name = 'χ'
CHI.plot()
CHI


finally we add the new coordinates to the image dataarray

In [ ]:
da = da_img.assign_coords(Q=Q,CHI=CHI)
da


we can use these multidimensional coordinates to extract an annulus

In [ ]:
da.where(((da.Q>0.15) & (da.Q<0.35))).plot(norm=LogNorm(1e-16,1e-4))


or a sector

In [ ]:
da.where((np.abs(da.CHI)>75) & (np.abs(da.CHI)<105)).plot(norm=LogNorm(1e-16,1e-4))


we can also use the multidimensional bins for groupby operations

In [ ]:
for i,sda in da.groupby_bins('Q',np.geomspace(1e-3,3.0,10),restore_coord_dims=True):
fig,ax = plt.subplots()
sda.unstack().sortby(['Qx','Qy']).plot(norm=LogNorm(1e-16,1e-4))
ax.set(title=i)

In [ ]:
group = (da
.groupby_bins('Q',np.geomspace(1e-3,3.0,10),restore_coord_dims=True)
.apply(lambda x: x.groupby_bins('CHI',np.arange(-180,180,5)).mean())
)
group.plot.line(x='CHI_bins')
plt.gca().get_legend().set_visible(False)  # The legend isn't helpful; just hide it.