arviz_stats.kde

Contents

arviz_stats.kde#

arviz_stats.kde(data, dim=None, group='posterior', var_names=None, filter_vars=None, coords=None, circular=False, **kwargs)[source]#

Compute the marginal kernel density estimates (KDE).

See the EABM chapter on Visualization of Random Variables with ArviZ for more details.

Parameters:
dataarray_like, xarray.DataArray, xarray.Dataset, xarray.DataTree, DataArrayGroupBy, DatasetGroupBy, or idata-like

Input data. It will have different pre-processing applied to it depending on its type:

  • array-like: call array layer within arviz-stats.

  • xarray object: apply dimension aware function to all relevant subsets

  • others: passed to arviz_base.convert_to_dataset then treated as xarray.Dataset. This option is discouraged due to needing this conversion which is completely automated and will be needed again in future executions or similar functions.

    It is recommended to first perform the conversion manually and then call arviz_stats.kde. This allows controlling the conversion step and inspecting its results.

dimsequence of hashable, optional

Dimensions to be reduced when computing the KDE. Default rcParams["data.sample_dims"].

grouphashable, default “posterior”

Group on which to compute the KDE

var_namesstr or list of str, optional

Names of the variables for which the KDE should be computed.

filter_vars{None, “like”, “regex”}, default None
coordsdict, optional

Dictionary of dimension/index names to coordinate values defining a subset of the data for which to perform the computation.

circularbool, default False
**kwargsany, optional

Forwarded to the array or dataarray interface for KDE.

Returns:
ndarray, xarray.DataArray, xarray.Dataset, xarray.DataTree

Requested KDE of the provided input. The xarray objects will have a kde_dim dimension and a plot_axis dimension with coordinates “x”, and “y”.

See also

arviz_stats.ecdf, arviz_stats.histogram, arviz_stats.qds

Alternative visual summaries for marginal distributions

arviz_plots.plot_dist

Examples

Calculate the KDE of a Normal random variable:

In [1]: import arviz_stats as azs
   ...: import numpy as np
   ...: data = np.random.default_rng().normal(size=2000)
   ...: azs.kde(data)
   ...: 
Out[1]: 
(array([-4.29998821, -4.28488109, -4.26977397, -4.25466685, -4.23955973,
        -4.22445261, -4.20934549, -4.19423837, -4.17913124, -4.16402412,
        -4.148917  , -4.13380988, -4.11870276, -4.10359564, -4.08848852,
        -4.0733814 , -4.05827427, -4.04316715, -4.02806003, -4.01295291,
        -3.99784579, -3.98273867, -3.96763155, -3.95252443, -3.9374173 ,
        -3.92231018, -3.90720306, -3.89209594, -3.87698882, -3.8618817 ,
        -3.84677458, -3.83166746, -3.81656033, -3.80145321, -3.78634609,
        -3.77123897, -3.75613185, -3.74102473, -3.72591761, -3.71081049,
        -3.69570336, -3.68059624, -3.66548912, -3.650382  , -3.63527488,
        -3.62016776, -3.60506064, -3.58995352, -3.57484639, -3.55973927,
        -3.54463215, -3.52952503, -3.51441791, -3.49931079, -3.48420367,
        -3.46909655, -3.45398943, -3.4388823 , -3.42377518, -3.40866806,
        -3.39356094, -3.37845382, -3.3633467 , -3.34823958, -3.33313246,
        -3.31802533, -3.30291821, -3.28781109, -3.27270397, -3.25759685,
        -3.24248973, -3.22738261, -3.21227549, -3.19716836, -3.18206124,
        -3.16695412, -3.151847  , -3.13673988, -3.12163276, -3.10652564,
        -3.09141852, -3.07631139, -3.06120427, -3.04609715, -3.03099003,
        -3.01588291, -3.00077579, -2.98566867, -2.97056155, -2.95545442,
        -2.9403473 , -2.92524018, -2.91013306, -2.89502594, -2.87991882,
        -2.8648117 , -2.84970458, -2.83459745, -2.81949033, -2.80438321,
        -2.78927609, -2.77416897, -2.75906185, -2.74395473, -2.72884761,
        -2.71374049, -2.69863336, -2.68352624, -2.66841912, -2.653312  ,
        -2.63820488, -2.62309776, -2.60799064, -2.59288352, -2.57777639,
        -2.56266927, -2.54756215, -2.53245503, -2.51734791, -2.50224079,
        -2.48713367, -2.47202655, -2.45691942, -2.4418123 , -2.42670518,
        -2.41159806, -2.39649094, -2.38138382, -2.3662767 , -2.35116958,
        -2.33606245, -2.32095533, -2.30584821, -2.29074109, -2.27563397,
        -2.26052685, -2.24541973, -2.23031261, -2.21520548, -2.20009836,
        -2.18499124, -2.16988412, -2.154777  , -2.13966988, -2.12456276,
        -2.10945564, -2.09434851, -2.07924139, -2.06413427, -2.04902715,
        -2.03392003, -2.01881291, -2.00370579, -1.98859867, -1.97349154,
        -1.95838442, -1.9432773 , -1.92817018, -1.91306306, -1.89795594,
        -1.88284882, -1.8677417 , -1.85263458, -1.83752745, -1.82242033,
        -1.80731321, -1.79220609, -1.77709897, -1.76199185, -1.74688473,
        -1.73177761, -1.71667048, -1.70156336, -1.68645624, -1.67134912,
        -1.656242  , -1.64113488, -1.62602776, -1.61092064, -1.59581351,
        -1.58070639, -1.56559927, -1.55049215, -1.53538503, -1.52027791,
        -1.50517079, -1.49006367, -1.47495654, -1.45984942, -1.4447423 ,
        -1.42963518, -1.41452806, -1.39942094, -1.38431382, -1.3692067 ,
        -1.35409957, -1.33899245, -1.32388533, -1.30877821, -1.29367109,
        -1.27856397, -1.26345685, -1.24834973, -1.2332426 , -1.21813548,
        -1.20302836, -1.18792124, -1.17281412, -1.157707  , -1.14259988,
        -1.12749276, -1.11238564, -1.09727851, -1.08217139, -1.06706427,
        -1.05195715, -1.03685003, -1.02174291, -1.00663579, -0.99152867,
        -0.97642154, -0.96131442, -0.9462073 , -0.93110018, -0.91599306,
        -0.90088594, -0.88577882, -0.8706717 , -0.85556457, -0.84045745,
        -0.82535033, -0.81024321, -0.79513609, -0.78002897, -0.76492185,
        -0.74981473, -0.7347076 , -0.71960048, -0.70449336, -0.68938624,
        -0.67427912, -0.659172  , -0.64406488, -0.62895776, -0.61385063,
        -0.59874351, -0.58363639, -0.56852927, -0.55342215, -0.53831503,
        -0.52320791, -0.50810079, -0.49299366, -0.47788654, -0.46277942,
        -0.4476723 , -0.43256518, -0.41745806, -0.40235094, -0.38724382,
        -0.37213669, -0.35702957, -0.34192245, -0.32681533, -0.31170821,
        -0.29660109, -0.28149397, -0.26638685, -0.25127973, -0.2361726 ,
        -0.22106548, -0.20595836, -0.19085124, -0.17574412, -0.160637  ,
        -0.14552988, -0.13042276, -0.11531563, -0.10020851, -0.08510139,
        -0.06999427, -0.05488715, -0.03978003, -0.02467291, -0.00956579,
         0.00554134,  0.02064846,  0.03575558,  0.0508627 ,  0.06596982,
         0.08107694,  0.09618406,  0.11129118,  0.12639831,  0.14150543,
         0.15661255,  0.17171967,  0.18682679,  0.20193391,  0.21704103,
         0.23214815,  0.24725528,  0.2623624 ,  0.27746952,  0.29257664,
         0.30768376,  0.32279088,  0.337898  ,  0.35300512,  0.36811225,
         0.38321937,  0.39832649,  0.41343361,  0.42854073,  0.44364785,
         0.45875497,  0.47386209,  0.48896921,  0.50407634,  0.51918346,
         0.53429058,  0.5493977 ,  0.56450482,  0.57961194,  0.59471906,
         0.60982618,  0.62493331,  0.64004043,  0.65514755,  0.67025467,
         0.68536179,  0.70046891,  0.71557603,  0.73068315,  0.74579028,
         0.7608974 ,  0.77600452,  0.79111164,  0.80621876,  0.82132588,
         0.836433  ,  0.85154012,  0.86664725,  0.88175437,  0.89686149,
         0.91196861,  0.92707573,  0.94218285,  0.95728997,  0.97239709,
         0.98750422,  1.00261134,  1.01771846,  1.03282558,  1.0479327 ,
         1.06303982,  1.07814694,  1.09325406,  1.10836119,  1.12346831,
         1.13857543,  1.15368255,  1.16878967,  1.18389679,  1.19900391,
         1.21411103,  1.22921816,  1.24432528,  1.2594324 ,  1.27453952,
         1.28964664,  1.30475376,  1.31986088,  1.334968  ,  1.35007512,
         1.36518225,  1.38028937,  1.39539649,  1.41050361,  1.42561073,
         1.44071785,  1.45582497,  1.47093209,  1.48603922,  1.50114634,
         1.51625346,  1.53136058,  1.5464677 ,  1.56157482,  1.57668194,
         1.59178906,  1.60689619,  1.62200331,  1.63711043,  1.65221755,
         1.66732467,  1.68243179,  1.69753891,  1.71264603,  1.72775316,
         1.74286028,  1.7579674 ,  1.77307452,  1.78818164,  1.80328876,
         1.81839588,  1.833503  ,  1.84861013,  1.86371725,  1.87882437,
         1.89393149,  1.90903861,  1.92414573,  1.93925285,  1.95435997,
         1.9694671 ,  1.98457422,  1.99968134,  2.01478846,  2.02989558,
         2.0450027 ,  2.06010982,  2.07521694,  2.09032406,  2.10543119,
         2.12053831,  2.13564543,  2.15075255,  2.16585967,  2.18096679,
         2.19607391,  2.21118103,  2.22628816,  2.24139528,  2.2565024 ,
         2.27160952,  2.28671664,  2.30182376,  2.31693088,  2.332038  ,
         2.34714513,  2.36225225,  2.37735937,  2.39246649,  2.40757361,
         2.42268073,  2.43778785,  2.45289497,  2.4680021 ,  2.48310922,
         2.49821634,  2.51332346,  2.52843058,  2.5435377 ,  2.55864482,
         2.57375194,  2.58885907,  2.60396619,  2.61907331,  2.63418043,
         2.64928755,  2.66439467,  2.67950179,  2.69460891,  2.70971604,
         2.72482316,  2.73993028,  2.7550374 ,  2.77014452,  2.78525164,
         2.80035876,  2.81546588,  2.83057301,  2.84568013,  2.86078725,
         2.87589437,  2.89100149,  2.90610861,  2.92121573,  2.93632285,
         2.95142997,  2.9665371 ,  2.98164422,  2.99675134,  3.01185846,
         3.02696558,  3.0420727 ,  3.05717982,  3.07228694,  3.08739407,
         3.10250119,  3.11760831,  3.13271543,  3.14782255,  3.16292967,
         3.17803679,  3.19314391,  3.20825104,  3.22335816,  3.23846528,
         3.2535724 ,  3.26867952,  3.28378664,  3.29889376,  3.31400088,
         3.32910801,  3.34421513,  3.35932225,  3.37442937,  3.38953649,
         3.40464361,  3.41975073]),
 array([1.95348357e-03, 1.94280916e-03, 1.92163501e-03, 1.89030570e-03,
        1.84932638e-03, 1.79934924e-03, 1.74115624e-03, 1.67563878e-03,
        1.60377501e-03, 1.52660563e-03, 1.44520899e-03, 1.36067632e-03,
        1.27408785e-03, 1.18649060e-03, 1.09887831e-03, 1.01217424e-03,
        9.27216884e-04, 8.44749135e-04, 7.65410743e-04, 6.89734140e-04,
        6.18143438e-04, 5.50956326e-04, 4.88388527e-04, 4.30803781e-04,
        3.77833542e-04, 3.29619440e-04, 2.86058256e-04, 2.46992579e-04,
        2.12222062e-04, 1.81514421e-04, 1.54859277e-04, 1.31589246e-04,
        1.11621612e-04, 9.47006888e-05, 8.05843710e-05, 6.92935485e-05,
        6.02285032e-05, 5.34058216e-05, 4.86960938e-05, 4.60060337e-05,
        4.52803901e-05, 4.65031834e-05, 4.96983125e-05, 5.51729229e-05,
        6.28715504e-05, 7.34492428e-05, 8.63146151e-05, 1.02095030e-04,
        1.21348231e-04, 1.43788843e-04, 1.70895933e-04, 2.02136626e-04,
        2.38132747e-04, 2.79594713e-04, 3.27092548e-04, 3.81054055e-04,
        4.42349136e-04, 5.11820891e-04, 5.89939759e-04, 6.77447011e-04,
        7.75352836e-04, 8.84532348e-04, 1.00566555e-03, 1.13915760e-03,
        1.28625006e-03, 1.44808931e-03, 1.62543674e-03, 1.81925178e-03,
        2.03026430e-03, 2.25985686e-03, 2.50882326e-03, 2.77812302e-03,
        3.06893101e-03, 3.38173168e-03, 3.71777482e-03, 4.07739810e-03,
        4.46117008e-03, 4.86998517e-03, 5.30398081e-03, 5.76327286e-03,
        6.24823286e-03, 6.75792610e-03, 7.29190703e-03, 7.84901246e-03,
        8.42793129e-03, 9.02725863e-03, 9.64557877e-03, 1.02798802e-02,
        1.09280444e-02, 1.15880793e-02, 1.22558565e-02, 1.29288646e-02,
        1.36039752e-02, 1.42780507e-02, 1.49480186e-02, 1.56114326e-02,
        1.62654994e-02, 1.69071900e-02, 1.75348226e-02, 1.81467635e-02,
        1.87420682e-02, 1.93201042e-02, 1.98802572e-02, 2.04232192e-02,
        2.09493087e-02, 2.14600843e-02, 2.19568766e-02, 2.24419727e-02,
        2.29165815e-02, 2.33827293e-02, 2.38434200e-02, 2.43004450e-02,
        2.47553502e-02, 2.52105424e-02, 2.56679004e-02, 2.61281450e-02,
        2.65924572e-02, 2.70610429e-02, 2.75353514e-02, 2.80145144e-02,
        2.84992969e-02, 2.89889219e-02, 2.94833317e-02, 2.99820924e-02,
        3.04858028e-02, 3.09945580e-02, 3.15077777e-02, 3.20285621e-02,
        3.25568786e-02, 3.30949034e-02, 3.36452105e-02, 3.42119328e-02,
        3.47981618e-02, 3.54084954e-02, 3.60476719e-02, 3.67209426e-02,
        3.74349963e-02, 3.81935072e-02, 3.90057643e-02, 3.98746146e-02,
        4.08055121e-02, 4.18047777e-02, 4.28775695e-02, 4.40260921e-02,
        4.52549154e-02, 4.65648547e-02, 4.79577974e-02, 4.94360166e-02,
        5.09973009e-02, 5.26409336e-02, 5.43631620e-02, 5.61630822e-02,
        5.80362281e-02, 5.99799876e-02, 6.19875755e-02, 6.40552710e-02,
        6.61801396e-02, 6.83578200e-02, 7.05828776e-02, 7.28505114e-02,
        7.51580725e-02, 7.75040815e-02, 7.98846721e-02, 8.22991267e-02,
        8.47446570e-02, 8.72218967e-02, 8.97296985e-02, 9.22681609e-02,
        9.48388508e-02, 9.74414732e-02, 1.00074028e-01, 1.02738742e-01,
        1.05434965e-01, 1.08161408e-01, 1.10919271e-01, 1.13705181e-01,
        1.16516475e-01, 1.19349443e-01, 1.22202530e-01, 1.25071041e-01,
        1.27951299e-01, 1.30834914e-01, 1.33719664e-01, 1.36596980e-01,
        1.39462675e-01, 1.42308803e-01, 1.45131244e-01, 1.47921544e-01,
        1.50675485e-01, 1.53388402e-01, 1.56053877e-01, 1.58667141e-01,
        1.61226047e-01, 1.63725694e-01, 1.66167873e-01, 1.68548744e-01,
        1.70870328e-01, 1.73132394e-01, 1.75336053e-01, 1.77483946e-01,
        1.79579924e-01, 1.81630215e-01, 1.83636657e-01, 1.85608135e-01,
        1.87549334e-01, 1.89466482e-01, 1.91369704e-01, 1.93263813e-01,
        1.95156372e-01, 1.97055378e-01, 1.98968045e-01, 2.00900611e-01,
        2.02860173e-01, 2.04852172e-01, 2.06881550e-01, 2.08953716e-01,
        2.11071755e-01, 2.13241022e-01, 2.15463139e-01, 2.17740831e-01,
        2.20073648e-01, 2.22465000e-01, 2.24915761e-01, 2.27427784e-01,
        2.30002374e-01, 2.32639134e-01, 2.35340556e-01, 2.38111117e-01,
        2.40953117e-01, 2.43867226e-01, 2.46862178e-01, 2.49939286e-01,
        2.53105519e-01, 2.56364049e-01, 2.59723265e-01, 2.63187046e-01,
        2.66760462e-01, 2.70450487e-01, 2.74260117e-01, 2.78190824e-01,
        2.82247145e-01, 2.86427438e-01, 2.90729968e-01, 2.95152011e-01,
        2.99688870e-01, 3.04335314e-01, 3.09081459e-01, 3.13915180e-01,
        3.18828062e-01, 3.23799433e-01, 3.28819627e-01, 3.33870756e-01,
        3.38933919e-01, 3.43988488e-01, 3.49017603e-01, 3.54000821e-01,
        3.58919534e-01, 3.63750890e-01, 3.68473685e-01, 3.73072499e-01,
        3.77526477e-01, 3.81819925e-01, 3.85932590e-01, 3.89849353e-01,
        3.93553177e-01, 3.97033258e-01, 4.00275902e-01, 4.03266954e-01,
        4.05999162e-01, 4.08465521e-01, 4.10657978e-01, 4.12569733e-01,
        4.14197265e-01, 4.15539001e-01, 4.16594999e-01, 4.17366042e-01,
        4.17855696e-01, 4.18066892e-01, 4.18006981e-01, 4.17683304e-01,
        4.17104808e-01, 4.16285373e-01, 4.15230441e-01, 4.13956993e-01,
        4.12476686e-01, 4.10804854e-01, 4.08958152e-01, 4.06953218e-01,
        4.04805905e-01, 4.02532659e-01, 4.00152261e-01, 3.97679470e-01,
        3.95136263e-01, 3.92535668e-01, 3.89896725e-01, 3.87233740e-01,
        3.84565076e-01, 3.81901968e-01, 3.79260662e-01, 3.76649516e-01,
        3.74079725e-01, 3.71560511e-01, 3.69099237e-01, 3.66699540e-01,
        3.64365134e-01, 3.62093645e-01, 3.59886572e-01, 3.57739350e-01,
        3.55645522e-01, 3.53601976e-01, 3.51599477e-01, 3.49626369e-01,
        3.47673299e-01, 3.45731563e-01, 3.43787318e-01, 3.41830749e-01,
        3.39850746e-01, 3.37834888e-01, 3.35775259e-01, 3.33663139e-01,
        3.31489953e-01, 3.29247969e-01, 3.26930624e-01, 3.24534019e-01,
        3.22052922e-01, 3.19486582e-01, 3.16831553e-01, 3.14085673e-01,
        3.11248277e-01, 3.08321312e-01, 3.05304409e-01, 3.02201974e-01,
        2.99012710e-01, 2.95738965e-01, 2.92382494e-01, 2.88947813e-01,
        2.85438951e-01, 2.81859226e-01, 2.78210302e-01, 2.74499547e-01,
        2.70733049e-01, 2.66914772e-01, 2.63052638e-01, 2.59153596e-01,
        2.55227564e-01, 2.51279713e-01, 2.47319830e-01, 2.43358251e-01,
        2.39402853e-01, 2.35463086e-01, 2.31552214e-01, 2.27675672e-01,
        2.23845629e-01, 2.20069253e-01, 2.16356610e-01, 2.12713277e-01,
        2.09145737e-01, 2.05661484e-01, 2.02266213e-01, 1.98962757e-01,
        1.95756800e-01, 1.92647711e-01, 1.89637978e-01, 1.86727376e-01,
        1.83916904e-01, 1.81203729e-01, 1.78583733e-01, 1.76053639e-01,
        1.73611072e-01, 1.71251118e-01, 1.68964750e-01, 1.66746519e-01,
        1.64588886e-01, 1.62484393e-01, 1.60425185e-01, 1.58401880e-01,
        1.56404844e-01, 1.54425665e-01, 1.52454835e-01, 1.50482542e-01,
        1.48499313e-01, 1.46498736e-01, 1.44473011e-01, 1.42412605e-01,
        1.40313122e-01, 1.38167406e-01, 1.35972582e-01, 1.33723463e-01,
        1.31420358e-01, 1.29061289e-01, 1.26646395e-01, 1.24178807e-01,
        1.21660912e-01, 1.19100154e-01, 1.16498934e-01, 1.13867073e-01,
        1.11208435e-01, 1.08532475e-01, 1.05847625e-01, 1.03164265e-01,
        1.00489286e-01, 9.78314670e-02, 9.52006625e-02, 9.26030571e-02,
        9.00470496e-02, 8.75377195e-02, 8.50807434e-02, 8.26788286e-02,
        8.03369021e-02, 7.80562158e-02, 7.58390190e-02, 7.36845579e-02,
        7.15924692e-02, 6.95606290e-02, 6.75870436e-02, 6.56692558e-02,
        6.38040299e-02, 6.19883836e-02, 6.02182982e-02, 5.84928689e-02,
        5.68064251e-02, 5.51567522e-02, 5.35427097e-02, 5.19608552e-02,
        5.04097814e-02, 4.88881519e-02, 4.73954020e-02, 4.59314270e-02,
        4.44951977e-02, 4.30882333e-02, 4.17103573e-02, 4.03615796e-02,
        3.90436434e-02, 3.77580896e-02, 3.65040990e-02, 3.52844378e-02,
        3.40972541e-02, 3.29446462e-02, 3.18277272e-02, 3.07447086e-02,
        2.96952817e-02, 2.86804050e-02, 2.76977376e-02, 2.67467787e-02,
        2.58274540e-02, 2.49387258e-02, 2.40776472e-02, 2.32437188e-02,
        2.24355495e-02, 2.16519414e-02, 2.08903885e-02, 2.01503370e-02,
        1.94293029e-02, 1.87277400e-02, 1.80434647e-02, 1.73754562e-02,
        1.67238437e-02, 1.60881709e-02, 1.54669350e-02, 1.48605124e-02,
        1.42686064e-02, 1.36909494e-02, 1.31285808e-02, 1.25810094e-02,
        1.20496866e-02, 1.15337334e-02, 1.10352284e-02, 1.05541787e-02,
        1.00912860e-02, 9.64701568e-03, 9.22192617e-03, 8.81746915e-03,
        8.43392577e-03, 8.07139829e-03, 7.73053753e-03, 7.41188079e-03,
        7.11507658e-03, 6.84074853e-03, 6.58884780e-03, 6.35866720e-03,
        6.14928731e-03, 5.96103237e-03, 5.79331163e-03, 5.64447241e-03,
        5.51421385e-03, 5.40087923e-03, 5.30431721e-03, 5.22260680e-03,
        5.15420717e-03, 5.09865851e-03, 5.05395937e-03, 5.01925417e-03,
        4.99289033e-03, 4.97391098e-03, 4.96118057e-03, 4.95381911e-03,
        4.95044511e-03, 4.95041387e-03, 4.95293141e-03, 4.95772051e-03,
        4.96360703e-03, 4.97009180e-03, 4.97734525e-03, 4.98420229e-03,
        4.99097792e-03, 4.99703779e-03, 5.00208526e-03, 5.00697434e-03,
        5.01089273e-03, 5.01450788e-03, 5.01638570e-03, 5.01745008e-03]),
 array(0.20394921))

Calculate the KDE for specific variables:

In [2]: import arviz_base as azb
   ...: dt = azb.load_arviz_data("centered_eight")
   ...: azs.kde(dt, var_names=["mu", "theta"])
   ...: 
Out[2]: 
<xarray.DataTree 'posterior'>
Group: /posterior
    Dimensions:    (plot_axis: 2, kde_dim: 512, school: 8)
    Coordinates:
      * plot_axis  (plot_axis) <U1 8B 'x' 'y'
        bw_mu      float64 8B 0.3569
      * school     (school) <U16 512B 'Choate' 'Deerfield' ... 'Mt. Hermon'
        bw_theta   (school) float64 64B 0.5168 0.5102 0.5383 ... 0.4794 0.5383
    Dimensions without coordinates: kde_dim
    Data variables:
        mu         (plot_axis, kde_dim) float64 8kB -7.485 -7.435 ... 0.001112
        theta      (plot_axis, school, kde_dim) float64 66kB -14.87 ... 0.0007317

Calculate the KDE also over the school dimension (for variables where present):

In [3]: azs.kde(dt, dim=["chain", "draw", "school"])
Out[3]: 
<xarray.DataTree 'posterior'>
Group: /posterior
    Dimensions:    (plot_axis: 2, kde_dim: 512)
    Coordinates:
      * plot_axis  (plot_axis) <U1 8B 'x' 'y'
        bw_mu      float64 8B 0.3569
        bw_theta   float64 8B 0.3822
        bw_tau     float64 8B 0.2737
    Dimensions without coordinates: kde_dim
    Data variables:
        mu         (plot_axis, kde_dim) float64 8kB -7.485 -7.435 ... 0.001112
        theta      (plot_axis, kde_dim) float64 8kB -29.63 -29.49 ... 0.0001257
        tau        (plot_axis, kde_dim) float64 8kB 0.9156 0.9539 ... 0.002818