From 62ca314bf66494e1035036ba40455206ff2a802c Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com> Date: Wed, 15 Jan 2025 18:34:38 -0600 Subject: [PATCH] Weighted Average (#833) * boilerplate * update boilerplate * fix boilerplate * add quad hexagon to tests * write tests, work on api design * asv benchmark * update asv benchmark * Committing weighted mean modifications from rtam/weighted mean to uxarray/weighted-mean (#866) * updated mean function with weighted arg * updated weighted-mean functionality in dataarray.py * edited weights to dask array --------- Co-authored-by: Rachel Yuen Sum Tam Co-authored-by: Rachel Yuen Sum Tam * some cleanup * fix tests * add initial dask test cases * use parametrize * add boilerplate for example in docstring * update docstring * added examples to weighted mean API, userguide and test case * cleaned userguide output * restarted kernel * removed duplicate notebook * run pre-commit * re-run notebook * Update uxarray/core/dataarray.py Co-authored-by: Aaron Zedwick <95507181+aaronzedwick@users.noreply.github.com> * Update uxarray/core/dataarray.py Co-authored-by: Aaron Zedwick <95507181+aaronzedwick@users.noreply.github.com> * update docstring, api * update toc * update notebook * update docstring * clean and restart notebook --------- Co-authored-by: Rachel Tam <116197442+rytam2@users.noreply.github.com> Co-authored-by: Rachel Yuen Sum Tam Co-authored-by: Rachel Yuen Sum Tam Co-authored-by: Rachel Tam Co-authored-by: Aaron Zedwick <95507181+aaronzedwick@users.noreply.github.com> --- .../weightedmean/edge_face_weightedmean.png | Bin 0 -> 26476 bytes docs/api.rst | 12 +- docs/user-guide/weighted_mean.ipynb | 208 ++++++++++++++++++ docs/userguide.rst | 4 + test/test_weighted_mean.py | 147 +++++++++++++ uxarray/core/dataarray.py | 71 ++++++ 6 files changed, 435 insertions(+), 7 deletions(-) create mode 100644 docs/_static/examples/weightedmean/edge_face_weightedmean.png create mode 100644 docs/user-guide/weighted_mean.ipynb create mode 100644 test/test_weighted_mean.py diff --git a/docs/_static/examples/weightedmean/edge_face_weightedmean.png b/docs/_static/examples/weightedmean/edge_face_weightedmean.png new file mode 100644 index 0000000000000000000000000000000000000000..868264662d9bc78a36f503808ec7fb347967cbb1 GIT binary patch literal 26476 zcmd43byQW`7e9JHBt!*4KtaNw1*BEF1nEXv1pzsP(hZ^*ARz)0igb53VgSO^R*?tZqA>i+nhQF>b2x?`hS$@#}1#Y^)jXOZ2RL-LiK#j_7sXj0`KETGi#uoEpof zGD_FXH8SU1)BU=!Gi3Ve`&k$a{KHI5*dzPrBZh=0&haOj5vk>W{)BY}kLvi7hsX7* zQ^%hOpWT2NBmXe{un3Mn2!H!9f^qx__cjs#@dp%!> zsxJM1I9?dg|4&!)|GkmbX4%elh51qPg_Coe#+6OmI=M2{Rp_hv{fRwC!4s-i?i75e ztCA>h+kCwEu(6i)7%p-&3v7LM-0U?TBt`x#%mv$~EMdF}omp;#O+aPlwmU;#NCKm@ zHhH0Idv^(k+NI^QUEk9L^>nSxr8}c?`*w#xm-+6gzbY#ya)7iKs=`O{Nko9Z~`%iHkvXy4enRe~&Rsq$z@ z067uVESfn#bzD|{O}-=$e5O50A?Q+6JQe9|R)%wYKa4mn{~9XihCFn46OXeRk<>dv zUB%!_(yz5@_z%|G;~S0ZsSJ9u9^pmcp)2EMK80hP!tO~Z{9rvZR9-44KO>Mb`kVeM z7yVb497=oo$43WS_WscA2=Ol$VmjHZ=Z`q(=njq z8L7>gAo3F1s4Gp=s6K?nsK3CxU@qiwH7>2^ko%EqAvTRJtSrXr4+5c3c*F(Yd*gzV zlyWs}EmPiV9i64T+d=new3R)hV%L=Kww{Ddw7SqN+)wT9a(0u-saeKoRT5tc-h!#* zioVjNiqd{Czl8id1DMX8DQfJxxw{*4#&>OJ4ebV=ORD4=$+hP^k+B}M?JhRPFg$}t zMN?1ZFoF@0_EyJ*Su32!Tq|QH#pt8Ad;{0hwx11QoHbp3(~tSBCuOw_l*2yhyYr6u#U zx3{+$KZQ9w5Lsp9q3H`>gtldQ=Ymos z>#nIF#Z%i@gQQ>0D>c-X{@MK&rM=(ALbNTin`?-lI$bGVW)?15YvNn@rSt9*WvoV~ zp3%s+7n_c)hsm-V$(p&Usj4Q85f~`pL&jI6J>isdaGOFB|CIH;G3uk~$X?c62T3Z& zxSge=L(BgjT{QmQ)K~Y#uJ`6c{bi2Ug$26sjO8MD%m=h?>NW@+ZW|OD$)>8O%jOz2 z4rt+Z79^2Mk)4GkHbZUJel8HWpr zU)@S}4wl@Px=?ZW@W^PoGIY@u2b(H$}eJ+ zA=|p=OJ1D+K$>gj_63sD{OQ_-cFvd9VAVr2D|YMB!6{vQEo%Bg_rfo~StsiIRx{}-<+)00nj9CV)G2{4ybI0`Q~G|0!>Jkb!4mF2G_Dua?o?-# zE}G)Bfq@X7I8%JO0NbLgNH|u~zhEwUq5sOL z4_n`hy-_94siHAIYjia>M{B+`Y|*!wgjx^L0Z-ND1&;RD?8+6s3fjF+zwUSYd&a5S zD_~thX!-EP43&uPsj664UpDT!;hWSu&)uzK@R7YFDwVdWBUzFhs%X@}|$8sNLhxCJhjP zHAX2_nm6)4n~DEa!ZN`a7UT0&qJw`cA;N-Vu*B9{`|(q|mv5-udgzRpVg!dYbZvvp z4?aqDXXrL2DaAGj?#+)CAO-=SC5t(+@(K7MqsfLa3!m(%Sm|1aBO$HV=O`vOu>x3; zr>R8Too!GPo*3y=cV0Eor@VmiQLH+rDYka0oq3_X3k*B+9M)p_)=*n7QBY5-EKZ8R zU2rSJdM&{WZC;w(QU|*moZx@boU40k$n+ed&*^ncY1IILVi5Og!PI8)3`)Wb(B}Tm z`~2*C44Uc^+c~qZb9byKdu%t_)tNs#s90g5JUY^~^6MHRc$;r(<<+{Zw@MoQ@+L8A z4Cl7Hi&m2Gw&idXF0b zU#Tix8zx_v^xZ-EA*Z=?quPM07BW_^&r^o4q64nfeflnePZjmeB#SRU!*M>lreSrW zeyr(xmgAaHsd$o?V%f}XFn6mR#&Xxw)r-m9Cw$l{Zc*@BW)V@A- zB~ZQM0vmgBv2>{*U-*JQEJ9r_15Kr31xmOx`D~}{Mx^uNeTP$#H+>()tSn=nzapiH z55Cju-RkL9sb#FsIjJA%6?nkr=ej;mjgjynO>o{?FTaVXgy@pE$PU!93Jv3 zk1a|wn=@5wW2_<$^PPK5VUaZ~mrgPwd~+0%i$g5%NbDt-T zm5GVtU~h}J+kn+d8AK&7wQ#j{s&}y7k<5JcO`Q0V?b;T|}gV zMqapcozwNbn7{y?p7Es?63jF2=dpkFGcQvww!50os7zw&=S4^1EG2jfyMWHPEW*_7 zlpI5wIOJ=j3*L9$!NjGiBa8DU<&V6kf={EuppT*v0A&7zoX0FQYwqVe`bZO;YH-FG z{ee_-UFn?0Uu;9dgxo9oiN2O7r#;C2e#K?E$^suXtW*)j6gU_3CJnB9pvD84-&D(qe2o(f zH^FZWXxUw{ZeS0Vf8KE0SooG=MO2DO)fiZbe>ak(t{$uCb!*TVEc--P z-CD8BOn+e}z>D|VVyv*+Ub)2yIB(jZCm6MN5hx$do!dX1QdHD45B|sK$Ln(zGLEzN zqN({%Lde!}`rH(G*5K4W1`D&OBs|Oj<_;_jVbjeX1A;AvijBGMQpw&=^;_5hj)I@q z-qqbRq&lWlXih_uz5l3og-h}O{bu5m>@dhnX#~O4~5!oJNh!f2JL$O z@=rs#9PgSgdr*VVQL|}Xm4cYG@w*p3tO8Z~FeUSAH#z}&V zrMoX5y%6;6o#>3&S;MaluP$+$cHBaDrLluM1K+aky?R;=poV`y^63D^AU0&>J`~TU zqz-8Pa@tzdoe|#Zj5nAQAZjhPT04z(>i@JH2V9p-;g}GIq|-+_#DuXMy#C_I^~Gk* zwoB7EY-eq%lw~DZRyX_c=`jNohUkuGO4Mn8$&+Ok8f@8L`+8)9CTTyX4B?l#C!em( zO-jzz#!u5%igz$Q!Fj2P;8Qi9jPCNg>N`GmHI(+D;EGG~Qlxs)h-RUq0rUIsVao&}>_!;SEJgVPB@ib!LIg9`g~GaT`;jl~-p(lSC{(dvlkUQt z;Dn0>zKLL?DUw=@x{G-3a2nK%a25*H8!KG`>=x}H5ua`ckrg168}%|Qej%@ zb3fX0j{pv!@uqeGOmCt^f?1erc@^wy*!sLyXS{5BLd96kSuB+lZV__fJH5U6DR6Lv zuWII-27ggnM$oBtAN<1TX5* zkJr7_b_s%a52Dm+&tB&EX^uh2rFM=K@iNv($ga&QAHkDeM`9-(0Dfyw{-;u=o?RxL|fhgo5+eCL42~6c|yZNkQb)=s5MoIGoqgLaWe1?(&)O zcQ$nPbuliCIyW3J5yWJISk|nvw01Z0V5b59@ieg*JX+o@|<@- zJz^17VM4?@{I)y59QbjXhqwr_xqq+hwGIss=NO`}V^AZ!quk>_NTG0XJ&{~M<%P_v zpCFVKMiJhGuiL&3=Zfz^*Zh+I{dyZAxgFmBHhJlS-{qJ5?ywLME z&i}^w(jRW^A^G1xu=Np=Ug4xCL4?q!qWNzYE(l+5b2q70T=lcJ;dQEc;?yV{Q-z>05VA!=p|Bdr%Eiz})laYuZU2~jtU5kU< z?B&VLT!1D)alHC(4QlC;=2J+sm+5Rrasr$oMmEceodiD7Lwqk$BZxqd_;3=q4dw$O zllBDP;)Bip#uT;Gfymk$@B}E3sTmNqWA@&OS>IFP12Ve524utc)V%wp#c@+J$Jz=t4LCCh5h5{LrutZQ+ z*zM;r2y)t)N&Gs{wF*$|X-4rR`P%MWW+{slJq&>v$P}s_AZZ&<5%52Z*uS=SeC-o} zCXeW3Lr1!p5^pqL_ve2+#urBb{ayyHO2G`+?e~{Cs|SFBw7}1K2P@*SJ=y%#Qaj^@ zmR7#$gx-frh>lrufmZ?A?6UFaMuTCxkX}3;<6|bFsVV)_X1b;7v-{!Bq!*bU10|c}RP@+P z(MLMD6fP0f5F-9JR8C3e`LSDq0FfJkf}De{e)Ln_)FPyoTjJKL)ja>A9>(M4)yd zq#g(feDOX}F2cw-R+%G-1Bvjj-+P~>J9s_lR>}6L5A&Biv6@cd-%ix!`W$vcj>Zu| z%2u-gs=3+9*8_Jlk9OI%%<#U{0B{)$X9svw5KoZe=WxyoPUxzhU>E9Fdz8qjwdEfK z5v;l;MJ%;9kYgBVJ^GG5MBVR~_FcPqldg2_sZE|^+@~OPkf+UQclWtkNAl?g*gj^; zlFt=>LjFaOT>o8l9&(gq3c3F!&MVY?s;gSrudCuE8oVXC8J=~GLtg&NU0XAW5y0f; zUeSJo`Dl&3&wiPt9pq>4=ptb(ClKv>4@m5hn6MY<@cFySl4Y+&#XNBw*{IPTb#SCV z%26f%wrmXFPvi7OHhDv`ZfSwzbo3iDS{)?0p?|cifpGj3V@dUaME9i6NHDaj2O(vhRZRo_|9R%u3E@5?)6qK#($9b*_& zUB6R#v~|=M_kp20^=6Zd*XX@+{Tq+sD$8|C?Katu2lw*}4W<`))PvdB^btw2Tef6k zlNZ2Z2WWcFbolC#Xh?t4aron9S2~)QS$Wte`&CdNK_XuZZP}OiTd=EACgf(ee>`Mg zH=T4s{F3XqX_X+{H6a#?>FWLPK>p&?j7tn95ZeikYm*j@q4vYBya}FJOzRcWHyC+H z4i$0k5X(_ZMQanns%yGfv1h?&Gzt`> z!TH-GsPVfayMcVOy0ApD#z;q$gN-6uZma8PfJ=>Z6Oi}if za#y)=UiI`y(etsobKl;N6J|ySd>9eYYnLc$G4%+zjBI07xfo5w&PJUTEiPM)wjxhS zO531fpqI^`d?DhNe#NkOcWc0UFfzCw0!bNzTth&1)dJFyOt`5WAu)%Wy}8B_xy^#5 zdB3(cy7et zs<+|(QYl?5N!9j=&)7>=ja-1_)S{9RWAPfKV#4uRi@w>lYG}9}n1cb4VI2Bt_Y;l6 zqry7NEnLfkZ2cUWF_-Z-4=fn@utRz)}Js`c>$_l20gRFW20N5NdM z<^2GEG zlmcZipgt>6ouee0Xg$?r*QG=`_V&8p+hmkkXG-u+``1ImkBohioN=aBXO)u)SKD#u^vLX)YMZ~K%y&{i3Q=$}m z@7GU{mFXTQf5S19?Un}GT$Or-oRAuiYCPX19YfRzfh_^f2CY;N&RxvF=H>4NA`n1FM^ zE{0l>4FKPTnMz4&SFXu{sphXsk@uUxSc?`8yE&U=mT!9BkvDugT;Xa0jCyoTOiY_O zTg7&avn?Lw*at}}JTJCs_Dq&ipyip?VF#6Hm~TO<(<_|oDbW`92MOqUx;$kJQlWUd zZaWk4HJy8ItH0?NHC68ozkEhWxK^y1XCkzq|8_Om1p`H`n1T8Gx}MCr37J39r;$zC zg}Pb}q2NF8IN-s#CyoOK<9Hhj0ssmIv>Yk~4Jb#SljdCjy#b&8Hj^REZ~itTS8JR!8)75P zfBT+9t|t52ObX2}PLMqf5K1UquKzGGMlYaT;fIJ3JU;wjgXi;$|;g9emWS&p4{tZS? zok3Ro{_#)*`$4608Avooa3X^t-^I$XFY>^W5*@FZpo-O06WH+tu zM!8XtW0Y1TJ`UCPN$)S}xF2lT9Y8$84Os*(;LUR9ygVGvdBCG!5|%p-a+kGx$NjiW zTE1ma+x0xMkb#i+6d~;1w=ctKkU5SNdI$0W1pYJw`iNwu*w_8xm2bOI;t*#9i5%Zz zJoaDBP+!E5Dp|RFrh-K8R%#u)OWva&*a&86t3(3~T)b>Z{F`!=d;0}irgO&03S<(% znfLWn#v?2xFaVqgTJ0@!%oYB7-3iB6rs2bDK+8YP}4{ z^1r$t_TB`F5Q(}bvumhwS#mb|T!hQ8vHB|yH`uH}N)#dJx*H1Mt=Al3nH~Rlk|EpF;p@YDTPuXJ&KyFNu1z)b0^P^tI|G}IpCU0hA2L?fz;E_u*ov6_mV$K~ zhEZC65>Jo~d%Pr|0?9S8j<51$pFsOm!m@msTLcSVSg=KI`V}5=V8y0vI2YsMIl?n9 zwHyv$tu}XqO_XG^5OQCG^t|p7h%wGNHA0f|4a<^pYA!*xOrRjdyZTp~+pZwN zrCZd3Sx3olc^%Ig(SSzeNNY9gP_icgL(W3OL z8;D(a5-7M#0yN8=@(&`LK+eQvIW(}y+^3&@Dc_P1f$-yN`4fN)j%5HKKC+gDb1k704Ya1pY zxZ1pOEvRS_qO@LsFXkgrsL_(S4C`jh&e#m;5{ndxIp)l=$A!sW^#~se@w;JZGV@jK z%jam44W|2ImdY~{TWcJ)|GY83rE~c}eGwFA&as&pogeBhQA!$Inzna>h&~}6-3cll zv&@L)mqxs5f%M8EiK_}M8F8Qzm;|NMaTdH6A`H2`&4q=nKkqr4vYP9W*p7I(0w}tx zzqio!Ak{J-r}v=dubZ!(LQXSbjQyh!vLs#z*p#w!m)hy@=e0i)_|wLs+7uDC7K zDiBdZ!mR$P1b*mLHchKT^9IC-SI!(sezItHp1Ts`y0=gm4cq@$8BfmirLh`+f&>;rbrMxPI_nCi zebK}XpG-CyfSC0m2}_^C-l<+z{UWM(I{_!y4Px+sUN8a5H~bVCHTSVDL8c4-#{oA_M9IObR$>FA=V}7l5eg4JoV6^yb@*Dje!p zexSo_bG0n`gS{uP3KU7obx8Jfhje7JtXw1`>vwLQzP$lY4c8ThY~ep{$X}%*MI9&z zsmX*ODoN_5xy(w7k60goOx@oCd|XpL)8;!25O+Y2)@Ci9T5sJAG5*e7`dIi$DKzvB zXYQE1P2ga%av2=o-y1%%Yo-|vJ`kWSW$;5JdIL(*;_m*}_L4@$x_>uD@V+r>`h3~6 zB}$}05SD<`sNrq#P>Sb7df`w`P0NxFyIPC^C|wwgsDJ6*)u4TB^B#|w?iIKsm0d(q zWj>Z~2ssTFF01j{(5)L)V_zrgKfM74PHoqnbcv55n&ZL!FSQlDZyQs+9nN(I`(P>L z$XwOeu4~KQD7@#FDDgx~?QvdQiHX4Sw^IdA4dRL&C1fC8srOZinYr=9%U7F|&SkWK zz-(KnhsXJhyOiB~kGiCtxO)M9ik!*IZw&n}91;{XdmX)8C}Jc*D9gtsa$Z6=`h%qn zh^8-*2PxMOerYFG;QwHwa)3+8b_J`aW;dc^Uhn0`wFm2iMe{lJpwZhX)F5NEwSJq> ztjXE>$}(_()PR&t&#Y3TA|{N@^ZlIC2IPkbHviG)QNBDYN!xcfMd%e>|SNEaK zmjR&bUxY*ZT_0(V+GigUjvLK{8p&ZID?Up3D&*m`~39 z2b9FmK=}39bDBLdsAtf!@&J~2ut8G4Io~~HjF#q-{GEx5Yg?5^gMM$4o?l>Q>JMts zFE9(!<}BD(>zyMW<-QmrT2U}<=})zM0Q1UWPodk{nN7_c=o&ZVs#0N~4#eD)Vo6Pk zuK~W&mRiwduc@oX;+hEzxOrkNhK^+|`={@nrrodaxwF|3?q%Zs)_xRZ&ISFKbixJEvH^ zTj3F-E4D}R_qIWx$PI3`lF~#?6?>Mhj{YZ_LDkE}TuC-%JXgzFve_X;x)^l#{oCxC zm*ti(*1d^K*np&Sc2%XR4L8$f#oppvq=cdp;7yNxC5G-;hmMqGm?P_#bBa%bES|ZP zCW~w$A!mVhl!XhkHdk;(1useKlJL9VIW|8im+VRXjYyPaN^+R^UXQZc^2rA4z4^Wj zw0gg^M(}S_+T^7`qd{*M9LU9*8?&#U2Id?CY|e~jr(?=zCcd#`BU`;ZlMmW5H;-NW zuZwBJ*3AM&f4+O!>Av&$s36!YcrR;)PeP6mL5s;s^|9$r2;4G zaS16bOI1q^k;}VxBLda6gt@Z>R2=7AccwkaPJ<@^ zJe!q{?gi4^&^vWzsxM}(JJjF1nQP@WIRS&+`vfXcwcu7P0u0nD9=C|_Gm^dyso#|6 zqSCyvZ22hIAHSkh)#Bw5ww>XF#E2wti1*$=DXJVY< zBUwq`>2@z^LXRo`bFC%$ZT;b7(9AatfX%RL9RqXPf98%pb|}5xDd~%K`4eD>{4b|I zpYO>o;C=g8n%2q2qK}j%IINg6eB>%hSP*JEa5F?pr~nbz^IT3}BR-em0%_e_+N5>K zCTEkpLu&1%b%~pAkDc9%AoL}UUka_*J0njL2V`|8k4UCenL$#FxLmi7kIw@8#brGR zf;M*?RHI|)7RiqkJ9PC_!r36>LnY=D_yrQ?1?e7P+~|;l@)>*Y?@>jv!<@nIKtDK< z#Rlm=e*{^M2}c6Bw@8n(T0xkrBjZ`T4nWpa=CNykbu>d)x|0rM3}D%%zBA-CqC(mLSg`q+oXkNa>L2;i0;EROs1XX_@sFQAi%28PZH6n2pzJp*Krn|*B-#HORITv?FBj|Ol$U?9Qf zEdv&7PA~ZLdjuAs2!W%d^E_>+6+q1I^-KV?CQ8LO)MC7_6#hs}E!~IY2U*E_o6Oi^ zszhZlI-@wjrhMK~1fr&X#X^h0$9iO%>$+XGOuGW3lwy_jhC~s&Ea@O&D^Hn~S~m2n zn?cBEofsXTx-L)($I}WnbB}Tx%Cr#iD>HP5xF78e2h%wsK2NApmAL;hwV@_waMH}9 zZJ`!(Mz;N{y^|UwN^8=#?s)D}pw0&8URsOjkRedZ4S#jYxnG?51N;$8^%KxXQx~2V zNo0K8sk1y_-JJh*2>~PDg?^TqCdhH>M_KK-2`D>mLXC^+ZTU+97PB#F&oONd@P=>n z(5>_(4A*cr^BdH-ra(3Ysy!Mn#!Rk)x+n7~82-V8;SUL$8uq^-^cVP(uT@!(r21Ju z{;zsC{ECR&K6h=A-#MYDMtu|(MO7hH&XUoK*|0yIq_?q!{wJny@oQ%!B zy6bS6ONo?a|Ekw__q;*2eQMaRQU&i+wqpoO#sg)DsKgqFe^4tkAU~j%ojG*6au@*H z`FEs3NsA|aVG0 zDLXHK4=TIPC4XL1hyrVD+dPR=*BNlI&z%rvQ0TL3oMY)IW337r4TzH4`-%YKpnhW8mMLe9nDufp@FH?hio|>jeeJ(&LOqqglV zLaG@)?Lkn`dl{db*$*(9$2{XU+9!v!XXUbjvIAO&p^jgfn3jvGPH~gUDss(=U;`)b5jY0M6gDRcmd>mP> zHf0Hm8+~p)kKXiZ;<_U83Z|c{@hzDwFHEXnJ4J*R4@1E86o*6Es1j6+cC0%eu1l5} z`@-qkzT>9g(DHO9H`>(`BB)R8=R=>fRT+>b`vfbeT8L@Nuhu2qSlJX9s@3;U4F zSHytY#hVNP(@<@&;1vC^${1LU#b*MTZIbU=k!m@BQiDx4y~Wm3p!#eO4f z&vBkWq@vgT=wR@=fe_j@yE~2h$r!h*amRvM|23}C5PyW=Wq<`n|0c!-PE7Zoa}f9? zTRbsS!vCLx(VOtu%P*lKx&@xE^wE%1ZmDjLyVPMI&-V43$9>E_i_)h$m zl>Y)Gw|PJ5+J%`DWlo0vSu6CR5`tiB5=?dWTIZ<_ERf(B5_MlkNy$Bbo=q8bi113E z^qi8@b^Ut)Aj@Xzl@IjUg=38Y+g7j=p7i+(r?KCHG8?3?WKtxa@Lb5msErqgw{XyNq*8qS;hqd@RIq&s0FPUau(~erfO+17V#?2wqcWGUp zX|GX_LyG-)?o-`gWuTh_M$%O<+^9sUfd}1nbFGd!GZgGxScU6ep8mzV?WB@{Y__lt&YyIV|HaS!jeb8X+3N10tgP$Kyc#Y+j?=# z1rDy=9+Fz7Z(9cJiJ}HA!$T$Fi^my-ZE?8 zw%33a-hPWkfcYXkY84^wg$tezh_VDjSZOc4^%Qb~PrK5>T(6{qU$qNqCuPwpVj-b- zbOB!bo5l?XLfj6lU&qvTZ$ejx=|-A_|JV+&VoWX>H`Q`a)Z|C<>y$s zw@7MXTkbxEuSQZ=ht~RSy=#)emaUwDtVmT0xL1%u$NJr^yGTM(oKNIE9JH-jVw2oI ziUT}3$YfrAj`{Z%mLNei0y&38qiXp09biV_s3Iz#(x0i$E>AiR1b1!8udhJ4#Vw@P zvOiE8a?nV0Gq|asXTj=~a2rrLO$jZtCPA7ki9;u+&boBL4Dw6vhHNCir2uU!%vX@2zc7wsNmf{Tj#I%z3`~hfsz~zxJte|Cjo}Z$#;@ZJt_s0I`f|;cq}| zfZ@;}E{C0q3#q!=%ch>Vq~HooD{+!p{}pMea475IzX6gvao2yJ8gks#18#c@Mu6Td z6i_G*h3+WB(Lzf5fmTG)Jm9j{XjhKZRrw^>_QRqLTsJt+zbvc*Ja}PWL3{H3CVZs<{+ErSAa6T1J77i0Q+@(1gr+ zUy{RX1fW^nuQ`fB?TDnulhp5jJz8^@Me5>;ULU7w?MfY1B1X#)CR)8=0}L5b&~upD zjdsuPARWZ@2dWW~wPY}ds7!o${LAdBp4eTa{Q0jyEI{*yLxA=18Sy}BE@VS*Ri&%R z-;adi&KrNeje(TIi;0r!^&bf`UDdvP@%t3CW~pqOruREfv!Bk);aoK3fX6v!fG z`C!Z6eLr>w0*-}j)pz!NPZPub{CLxU0Mzrtv1qAPdJ2er^8C&mLZwh3Fy)wCbC~ zV8Pj8bI=t^GUUP1D;6oT2*`u%Sq$Q1yJn*|ht)_bdh4H5Gy|3EdTVGU@SHn~mxB%} z@o;`cKdBG#LMoXm_g7`6fs;H7$4W208RD-S>rhf%q+u+t2o*iJ-ZVj}B^A?miSHzW z7L6tENj>S7S#7l*5Et)?S)aZmtCW~_N|URcX z`2k3f42>Bw8xezN6vL+)7vC6XqBXfpn4O6q2p{0A{(aCX6bd> zoPXMU4&^c1Q}o;*zWa8o?J`Jsv;4Vh-9am5)}8uR*RJ=M4v2g2`zAA#Sj?3mprdW< z-YZMDMkIlxtuq(?>D*6r@dwM5hGe-nUBl$*QnK^dVnDTd?~$|IT?F|e(b|E$A&}QV zxjWSs^eLt0+AZ{XRe+deYu*l+cNW*L0Mfj`4RUVDP9B$l1lL>HLTzC*MCM*lvZDRR zmdWMi6$CO%Hv_ZXl+AJxD#b6hOMI;iap-OVlR|f~0*`f2l6^QzdUYCj3knq}}FK9wO6UL=3e9Q7)&$3a=QkWyTkW{}_UJh#*fejj$% z)>QdEmF*VPf$Ck&6kP{q)5_Thr%KeS?E+1Zz`ZB+2oU^8e}TlchDi$>nD6$}iWrih z@flsDD4*VT`f5XW$(9PDCT3TEe73R{Y?-ioQ6Q!aPPJ0|O{sQcTEON@3d8T%IixN| z*GaHs(RJk)*<$bNx=tB7HbVuGH1?;y%J2bFh5#Y)hAtUXfG~ovQtjs5s@QSB8Zz~? zu3+^)PRDrVwLQ$7rHZF3A>n?fy7jSEbutHFS{V3%_H^?^qrkyZdp2Hfr?SZkNc3<) z4Qa}qxwKr698Ly=-`fS7&BW_qP1o8`ILnpq8s}(UM(c8c1IT}CX9^^f`txj5`pnPH z4vpdLX91?hQ;OYXxZJ1wBR?>?yx>KoCNl-gr@QU8+R^!C8(rF)p%+Z6DV|DaXkOr& zn1FhV1t&xcsPdTGmH-EHvj_Ba`OP!i;R088o%MB*f{ie6X0B;VB^7(l^Vx57aZphl zB~-eO6Ty>!Hwx!U>B~l z#jrG7e6xFO`JB^0$T^VIyU*X5+uef7;~tA@l7$H7pXg4p{Sb1kWts?+s+W0DbZ^6> zbXum;o3&n=qn|Azd&z=8KdBMo2J`P2VVYSFUkuWb%UVLvvzdLUn`8LS^(Oe<+^;g4 zcPGJY&U*v`jee6-Rr%DE1i{(l@M~WbxLlhd7=RTrm44I06kO^Z(2Nag{eI4MDW&b_x;pr+9&= z!OJEbiWd(Wn6>mA{+y<`c@N}+&1T<{m#X0qoA2?@iDr%nbMulDGhK#^zjH%S=DUAY zkC;Q=)d?x@|3~)Nyp6~n;gua1KQf#~(L_QK2(YA0I*=i>8wi}f(G1uaeGf|A5kYrB zP)pDxoS__&XWq|)wm%3R%qUD`^aAn=x%1P~Lr7{lsty5vwcPphZNhzL&WOvos*NEP zbfMtne`0^mH|@-6erIZEH)yL-NGhXyM&+p+h2_I#S;dFzkQVEEs9kyI9ClKrT&1t7 zX&NZ%89T@QMp<g()IESBOM3i zMK0ml!<7bpcdDCG=%oF`{TJe1NsWKcc!=(OXlfsIe4g1)z1j4$Be}2f5eRb#)Pv-zPqt<}~qzaAg&;ogByo z3q(9q%}QZ^cI`!mU^h6YHbGL`FUo`g@*ol^S@4b@_ro-mo9j)Gqp);?ms7a^^69pb zyH8hq9VyE#9CAVrh&Ljyjwo%}pVL3$07IKz?&bzkK06m;+rq2GFYCsg_p!%SxAA@X zKmMLsRu0qRbg3$AC%(5tah`NkeyR-AoGqYuCS_t`^6sq!4;KyGAv0P{FVpF$Z})KJLj&A7;2VzAQjF@{FCU#u9AH7Js2b6aVz3c|Zpy&5t^g zl%jw6koAwzp^;K;2tKOI>|Ro6jqN6%^sP(!Aj6@4GDBcS+?-Ucw2svDI+Q6=i6#w)BysL11x zZpP>=(&^;}S=4BFl?U_U;ge6G<3YIGMLo)#;Z`t`KL-ovW40V%O$V8o6W2<>9*j{X z=71myg{6>Cr=}|t32z3H{>YzPjlP>;5tnOgt)=!89VP*m6JdEvOs+mPoqbU7E_L@| zCB&W<>EbDkF0oH_F8DUaK6mVJv4YUju6dlQ(d8hvx2I!Y$2E5hbdu!eJUD((`NYBxzqAY#9hnHJItQfe=KnBNjPg`HB# zi$bZzqd%8?L(6_%Zp`b~YCDkNBw>6lL#u~B%-h4d6;3FY*}dK+cuj9jzoOdGgG_pr zzm~oOEB>pgT@s#j_PzD<0p)PYGbCDFpasPkn)>yP#Anl^THp9k)!w#_)EqT z7nd_x5kA4>l5L|Mh@MQH3=6+1b-iZ`qB)n1P7MnJp_T@C9D}8eh~1)QT=+o{a08B@%M~ z?db^SEdTe`o*?|RdW~#dq^uA{cm;k$pUObvVR0GE!frw}>M#;K zAq~Co*_~|3QqmsAv4$)e5hE&PNre_mqL3xaq@zNX zvL#gfK9|<#oZs)i-{X&7dYOCXy6^kCujTXoeBL8hY6tQheEnQhrAI=X9KR(7Fu#Ei z8b06x#KA|GQ~BGFp(t#&B!Epwhc4QskpK5Q@)~Dy-GV9-i~aLW>4A9wqWn1e(hP0i z_dQY9a&ByxezDOwQ`iLo(f4&+N!Uw|U{q;R0&9Y~_Hqgs7mVSEEeLuD&3!kFW zfCov-}&mDq=a@n6k&yF75u8H2hBbDjAEssw_yBNsoj= zG73Pg%AYAKd6#2KYS-;`hACTk#M{HoL2J_>op8kklt zmLr%-B~Spbgmr@t1gat6QL*1-M~|51yCjhPdKEDTM32=+?SnAVi5b!hBmxGq{$ZAv zRID4Yfe9}h$C^?etL5xbzyO3dL^EP5pq1J_nX5K1wyUVS4FPS=6aOy&m8Q#Pbm+WY zc@%<7AY_!p?_?DJwQKyC@vH#lpKpF=G(RB|9EK^vVYu;P5V`EHw3;Y86$SC@ z=cLR!8gpL8&>0yXXX$0@IRcOVov*N|XDH52<@)63kHBa|bgOzk1_EUg=oU}VFDkIl zZA%%8%%GJATzXjNE8epstUG?j7~X|aw;vikxDeO@$%Cl_bv6ZwdJ~D-n;aF=U&#{R zISIQ+NGr6Xx9#Tj3^Rp~d-I+fyt9hc?BW&6N7F@@d#k8_gUk3pXWWYSb#A$FAT&p| z)|ZrQ817~F?sarA+wE-pWOnmMg;c1ZFN=K##{6Hzk<(SnrZF6*rl#__)ACoU+2ktn z+!Xv18lr<2h?$@_2h@N1H2j2|8&qes;oLhOvEiRSmpri3bQG+yFmL{J9MD*ca2zrG z{PQ?08q~OYny$QCRGRNM?)IN-ly5DiyqX4v#F6Fj*1}+VCfA3~2j$<2e=8rKF@{&B?_Wae#dRAAuHC z!EwwsAMI9s|2*m&Af!PZ?fH)QX@U&w^bgv|W%pSxWdGhL4r31hKFdcGzN;(3ma)w{ z-`p8V#QQ#`d{P6NgEd$GqrW#1nTdZ4l_Mr|>2>Re{Q1$MC{JumGkK(=G>A5hNwmF<4NfOX;#mY;fx`o1xndf;10W;q1{vu;Swh;%@`NO3 zk!cF&b~Oy}ZRQhi5L46i$R>;l(j?~3w2)%5mua5lkNVCU$@l_fDjeVNV^A?n%Qn3h zI4BoD0UFAvCa`?j_iL5pVmAq$sxbHLo0%CH^OWu;$dG}h8~7y#`N*$;lKtNRB^AM% zFXOdtk-}`rQ4Y3wL;Cq%DDJ(DtIEA?2~A3D9tHvKZ-YRr91P~?<9D8PiCH2vF{2HRBQ4ghMIJjhC{yl&c4UZV8GLV6J4%;S%;0QuxBrE~HOBy4ypL@V8>+#8 ztA^I&N>p`RSsSn@_atdIzIqrzgD_5LSm;f4goXv%s|X80m2lVX(lA~m556oXuBuMP zpI;i|T^^j@9q-|??Kx(%&DdO>{x}tDSKB)>VHoLO%o4dTL;#rG&=`1f0jrsm_U+F# zNyq*%8ItR&=^BQRz7{B5u?MjKER$|`yVh+yT!k8x3&a{a=F0u~0RQQb zIQyC-lfy~<2gj!KCt(XvRq4|c8FiYV1lERH=G$#=5?L(=0p)eCqG}gjxQz-5Vaj%vFFvxC0Porq z#I5qMHWST1$s#qHeEZ=J+Vjw4h7 z)gX@xe(I>{{c@ZtZ4J2bXxbR20p;mxRN4%e6@R1rE{HcRuh~<&<3YMObwqB_;jz8v z?J|PHmE$xllzSk%o*PGJ>Gu-dkKUtnzyqO+mSNh7nO)F~l1M<+`H^U@H}In(;8ylF z2CoNfl)cOly@q1_n2e$*f778k(Ql<1uSm{r>LH7c3_3@Vp7;x7VQuId*8c`H=J*l6 z*%tB$mFiD>< zIJEB2khO(UB4}gDoB34N-hjijWk^40k>mNqr<+>N$*&!*EOX9~GKzi1oWAOtCa}2D zfO2YGHKpUZZ_mY(L#clNds9H@7+#~XM4=MAkRj!h#YmlY6qjkm)pwY5;Zn_Kp0Dpv&~z# z3nlbo$mx><{bLRbfYdj`-~i;)Y7Cm7F;feny$4jK1R0k!WSTtT)n&cqU2T$mgClyhzqN_5*YbKRv1zft zrD<17cYf;LvXyY1$T{8JO`==<;o=0*(Z8Sp$?5AD&f^jeK5zDT+t%C$=YUJ-P<_7x z7FT)N(cQxLPw^?^sl2UP=5njw_SRuGYV5Gi-g|wrzcKY41I+<-d zs?t-bW<3t`WCsO8Z7=ILWfS9kAQ#OycGaD1)<}MjO2e;IgS6ck4&KhEC;E??ZDcD! zW6=?mwZa^xWPE1pwXnj<;%VR2p&fB{qYuy8rqoS?x`|W=`fWxaHOp}p{Ow~bC^@Hgi_FpGTF}wZ=mbpiUYf4v&&K)Gq%x;* zbZ-{|@&?6n-7reFH}~hY;xS;SgqKJZ8u4R+T4-3Wf3uG%(;H?iHy5EpbLzgd?GY02 zG_NP*@jR_+0KI9b(zjh7B&KHD%E*;2S=N2`A>cn+wE0?GqQRry?rU6+e7>Y+R*o~X z4R?I%7A7q9uc?{60!kS*gIgFP_KSO6^_8z;G12NtX@2+5&^e^w=IY?6p)t~`xJ z_JGk68bdG7v2mV>rCv+IA~llj{zZe{OMpViYmk^GSs22-B+~(mHH&uk0dUX+8Dp_g zFqB`#!Au}I+xm@@`g8@9dqBmj&$PAN@{HscTYSiZrI7yw#8tWWx^b6=gc<*inz4{O zQmHMRf&qPq;aZ#J(EOJLg|A`8B@%~+guNz*kmH2CxJu-h?l0C*=<|FHv|x8WCYWgv zX{%8u7fHQ2InmocJSaI^_ruDDr4LcvbyjE-*N%>SoPJTFl9-u*iVr2Wcp7CTW!uS_ zp_|^<8?eo{|BJL^|IV(bhqp8>H`4UK9rcR&D<6eY!6{Lm z9uvM5sFKhSajn7U@Pjxl=BVHSmnlGU-56YU_K5E3=clSzSw@YdpPsp`2*e^P4B>6j zrg#tvY2tu8t92k;l#_LYzgF#V&$&COmlQ~!QFMA%)40Yc`jy@sPdEYvw&9+?d;Rj2 zFnlXA&h|1`23(6`gc+-|27+8XMcF=*6P}bh$c%H=IGoIsvqN%VB{9#n-Y{v6ysDQu zpXY@M7PWvqJ8a0zHicu`)jq;f+-q%hTG_kxSn@eL8`=m;aW7P=!j~>|7JP#}_Z-Q3 zOULl%eQ*BK4q{U4rwaUJ#x7AiQXELYl}Un?uY@T?_N*Ye+Lq?o9abt4W*xcvnh_K4 z6IRBgl6S1jKJ6IlSIv#`K06g9!&8!SOAOtYM?xWbY#a=nQ<}3D1^x zN6;w)Py;WN1fzJ7_XFju^q!5Bx7E^m4ar?eu6Vg$*(>3B1_uv`2QQ8K!$ydwL_~G& z7LVR!>oRjrqLHf|)%%czaiKXapKzQA?StU%{C5gjR>xhCb7wk6)LX53UJ`)w2lZT{ zobp159ECHaQJpI>hV%=sXL;(MTTOep9w6Nhm*ubn%Dn>VR&9v%TMAK)1!G{6Xd{`G zsGYCp*J9gBzMRCpWX7xl^2H#P^{_*`k;szqm~&tsU+nIS0!004eUHphw3;r~d(`q8UT zG?yi_aj=xw0R&XS1_em*zSaEHpuka5f+4RTn1jo?4P@M}2jBZ7dt^MQe`gXmjUi`K zDOje%$tC#huV^UQo#Ff_lu)_#GXt!V_$U9N%wAV$_vN?$(wD)gDRfS8Jip!Lzv=|v QcJhhawRUP|YwQX9A5wF9@c;k- literal 0 HcmV?d00001 diff --git a/docs/api.rst b/docs/api.rst index 10307a16f..c066dac8c 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -387,6 +387,7 @@ Dual Mesh Construction Aggregations ------------ + Topological ~~~~~~~~~~~ @@ -413,16 +414,13 @@ on each face. UxDataArray.topological_all UxDataArray.topological_any - - -Intersections -~~~~~~~~~~~~~ - +Weighted +~~~~~~~~ .. autosummary:: :toctree: generated/ - grid.intersections.gca_gca_intersection - grid.intersections.gca_const_lat_intersection + UxDataArray.weighted_mean + Spherical Geometry diff --git a/docs/user-guide/weighted_mean.ipynb b/docs/user-guide/weighted_mean.ipynb new file mode 100644 index 000000000..4b6ec330c --- /dev/null +++ b/docs/user-guide/weighted_mean.ipynb @@ -0,0 +1,208 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "23f93b0f-62ac-4efd-975c-368477c9ebbe", + "metadata": {}, + "source": [ + "# Weighted Mean\n", + "\n", + "UXarray supports computing weighted means based on geometric properties such as face areas or edge lengths. In this section of the user guide, we explain how to calculate variable means using unstructured grid geometries as weights. The examples below demonstrate the application of weighted means using grid face areas and grid edge lengths as weights.\n" + ] + }, + { + "cell_type": "code", + "id": "ed637e46-672c-45ba-a5f1-ed3c3248c8a1", + "metadata": {}, + "source": [ + "import uxarray as ux\n", + "import xarray as xr\n", + "import cartopy.crs as ccrs\n", + "\n", + "import warnings\n", + "\n", + "warnings.filterwarnings(\"ignore\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "8518335d-2f75-4699-bc3d-862ab25830d1", + "metadata": {}, + "source": [ + "## Data\n", + "\n", + "Data used in this section is a 3-random-variable dataset on a quad hexagon mesh mapped to the edges and faces." + ] + }, + { + "cell_type": "code", + "id": "a2f58f07d7881d4f", + "metadata": {}, + "source": [ + "grid_path = \"../../test/meshfiles/ugrid/quad-hexagon/grid.nc\"\n", + "quad_hex_data_path_edge_centered = (\n", + " \"../../test/meshfiles/ugrid/quad-hexagon/random-edge-data.nc\"\n", + ")\n", + "quad_hex_data_path_face_centered = (\n", + " \"../../test/meshfiles/ugrid/quad-hexagon/random-face-data.nc\"\n", + ")\n", + "\n", + "uxds_face = ux.open_mfdataset(grid_path, quad_hex_data_path_face_centered)\n", + "uxds_edge = ux.open_mfdataset(grid_path, quad_hex_data_path_edge_centered)" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "id": "6fac99462efcccdc", + "metadata": {}, + "source": [ + "(\n", + " uxds_face.uxgrid.plot(line_color=\"black\")\n", + " * uxds_edge[\"random_data_edge\"]\n", + " .plot.points(\n", + " cmap=\"inferno\", size=150, marker=\"square\", clabel=None, tools=[\"hover\"]\n", + " )\n", + " .relabel(\"Edge Data\")\n", + " * uxds_face[\"random_data_face\"]\n", + " .plot.points(\n", + " cmap=\"inferno\", size=150, marker=\"triangle\", clabel=None, tools=[\"hover\"]\n", + " )\n", + " .relabel(\"Face Data\")\n", + ").opts(legend_position=\"top_right\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "bb77e5067497a0f5", + "metadata": {}, + "source": [ + "## Weighted Mean based on Face Areas\n", + "\n", + "Here we first look at the data values on each face and the faces' respective areas. " + ] + }, + { + "cell_type": "code", + "id": "c87fb018ca75ba5e", + "metadata": {}, + "source": [ + "uxds_face[\"random_data_face\"].values" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "def2b952ba7f1740", + "metadata": {}, + "source": "We can simply call `UxDataArray.weighted_mean()` on the UXDataArray to compute the weighted mean. The differences between the weighted mean and the regular mean is small since the area differences across the faces are small. " + }, + { + "cell_type": "code", + "id": "bd76b23993d9967f", + "metadata": {}, + "source": [ + "result = uxds_face[\"random_data_face\"].weighted_mean()\n", + "result.values" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "id": "a92f056d6b2a567b", + "metadata": {}, + "source": [ + "unweighted_result = uxds_face[\"random_data_face\"].mean()\n", + "unweighted_result.values" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "e8ddb288f83486b8", + "metadata": {}, + "source": [ + "## Weighted Mean Based on Edge Length\n", + "\n", + "Here we show the similar steps but for edge-centered data and the edge lengths. " + ] + }, + { + "cell_type": "code", + "id": "1b01581da7ebacc", + "metadata": {}, + "source": [ + "uxds_edge[\"random_data_edge\"].values" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "id": "a50c11e691337cff", + "metadata": {}, + "source": [ + "uxds_edge.uxgrid.edge_node_distances.data" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "647630d11f001fcf", + "metadata": {}, + "source": "The differences between weighted and unweighted mean is more drastic (~0.1 value difference) since the edge lengths have a larger variance. " + }, + { + "cell_type": "code", + "id": "6bace8386234b942", + "metadata": {}, + "source": [ + "result = uxds_edge[\"random_data_edge\"].weighted_mean()\n", + "result.values" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "id": "d81bc60c195be6ec", + "metadata": {}, + "source": [ + "unweighted_result = uxds_edge[\"random_data_edge\"].mean()\n", + "unweighted_result.values" + ], + "outputs": [], + "execution_count": null + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "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.10.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/userguide.rst b/docs/userguide.rst index b4e6249a0..b62cce86c 100644 --- a/docs/userguide.rst +++ b/docs/userguide.rst @@ -55,6 +55,9 @@ These user guides provide detailed explanations of the core functionality in UXa `Topological Aggregations `_ Aggregate data across grid dimensions +`Weighted Mean `_ + Compute the weighted average + `Calculus Operators `_ Apply calculus operators (gradient, integral) on unstructured grid data @@ -100,6 +103,7 @@ These user guides provide additional details about specific features in UXarray. user-guide/cross-sections.ipynb user-guide/remapping.ipynb user-guide/topological-aggregations.ipynb + user-guide/weighted_mean.ipynb user-guide/calculus.ipynb user-guide/tree_structures.ipynb user-guide/area_calc.ipynb diff --git a/test/test_weighted_mean.py b/test/test_weighted_mean.py new file mode 100644 index 000000000..f08b1c2a8 --- /dev/null +++ b/test/test_weighted_mean.py @@ -0,0 +1,147 @@ +import os +from pathlib import Path + +import pytest + +import dask.array as da +import numpy as np +import numpy.testing as nt + +import uxarray as ux + +current_path = Path(os.path.dirname(os.path.realpath(__file__))) + + +csne30_grid_path = current_path / 'meshfiles' / "ugrid" / "outCSne30" / "outCSne30.ug" +csne30_data_path = current_path / 'meshfiles' / "ugrid" / "outCSne30" / "outCSne30_vortex.nc" + +quad_hex_grid_path = current_path / 'meshfiles' / "ugrid" / "quad-hexagon" / "grid.nc" +quad_hex_data_path_face_centered = current_path / 'meshfiles' / "ugrid" / "quad-hexagon" / "data.nc" +quad_hex_data_path_edge_centered = current_path / 'meshfiles' / "ugrid" / "quad-hexagon" / "random-edge-data.nc" + + +def test_quad_hex_face_centered(): + """Compares the weighted average computation for the quad hexagon grid + using a face centered data variable to the expected value computed by + hand.""" + uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path_face_centered) + + # expected weighted average computed by hand + expected_weighted_mean = 297.55 + + # compute the weighted mean + result = uxds['t2m'].weighted_mean() + + # ensure values are within 3 decimal points of each other + nt.assert_almost_equal(result.values, expected_weighted_mean, decimal=3) + +def test_quad_hex_face_centered_dask(): + """Compares the weighted average computation for the quad hexagon grid + using a face centered data variable on a dask-backed UxDataset & Grid to the expected value computed by + hand.""" + uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path_face_centered) + + # data to be dask + uxda = uxds['t2m'].chunk(n_face=1) + + # weights to be dask + uxda.uxgrid.face_areas = uxda.uxgrid.face_areas.chunk(n_face=1) + + # create lazy result + lazy_result = uxda.weighted_mean() + + assert isinstance(lazy_result.data, da.Array) + + # compute result + computed_result = lazy_result.compute() + + assert isinstance(computed_result.data, np.ndarray) + + expected_weighted_mean = 297.55 + + # ensure values are within 3 decimal points of each other + nt.assert_almost_equal(computed_result.values, expected_weighted_mean, decimal=3) + +def test_quad_hex_edge_centered(): + """Compares the weighted average computation for the quad hexagon grid + using an edge centered data variable to the expected value computed by + hand.""" + uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path_edge_centered) + + # expected weighted average computed by hand + expected_weighted_mean = (uxds['random_data_edge'].values * uxds.uxgrid.edge_node_distances).sum() / uxds.uxgrid.edge_node_distances.sum() + + # compute the weighted mean + result = uxds['random_data_edge'].weighted_mean() + + nt.assert_equal(result, expected_weighted_mean) + +def test_quad_hex_edge_centered_dask(): + """Compares the weighted average computation for the quad hexagon grid + using an edge centered data variable on a dask-backed UxDataset & Grid to the expected value computed by + hand.""" + uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path_edge_centered) + + # data to be dask + uxda = uxds['random_data_edge'].chunk(n_edge=1) + + # weights to be dask + uxda.uxgrid.edge_node_distances = uxda.uxgrid.edge_node_distances.chunk(n_edge=1) + + # create lazy result + lazy_result = uxds['random_data_edge'].weighted_mean() + + assert isinstance(lazy_result.data, da.Array) + + # compute result + computed_result = lazy_result.compute() + + assert isinstance(computed_result.data, np.ndarray) + + # expected weighted average computed by hand + expected_weighted_mean = (uxds[ + 'random_data_edge'].values * uxds.uxgrid.edge_node_distances).sum() / uxds.uxgrid.edge_node_distances.sum() + + # ensure values are within 3 decimal points of each other + nt.assert_almost_equal(computed_result.values, expected_weighted_mean, decimal=3) + +def test_csne30_equal_area(): + """Compute the weighted average with a grid that has equal-area faces and + compare the result to the regular mean.""" + uxds = ux.open_dataset(csne30_grid_path, csne30_data_path) + face_areas = uxds.uxgrid.face_areas + + # set the area of each face to be one + uxds.uxgrid._ds['face_areas'].data = np.ones(uxds.uxgrid.n_face) + + + weighted_mean = uxds['psi'].weighted_mean() + unweighted_mean = uxds['psi'].mean() + + # with equal area, both should be equal + nt.assert_equal(weighted_mean, unweighted_mean) + +@pytest.mark.parametrize("chunk_size", [1, 2, 4]) +def test_csne30_equal_area_dask(chunk_size): + """Compares the weighted average computation for the quad hexagon grid + using a face centered data variable on a dask-backed UxDataset & Grid to the expected value computed by + hand.""" + uxds = ux.open_dataset(csne30_grid_path, csne30_data_path) + + # data and weights to be dask + uxda = uxds['psi'].chunk(n_face=chunk_size) + uxda.uxgrid.face_areas = uxda.uxgrid.face_areas.chunk(n_face=chunk_size) + + # Calculate lazy result + lazy_result = uxds['psi'].weighted_mean() + assert isinstance(lazy_result.data, da.Array) + + # compute result + computed_result = lazy_result.compute() + assert isinstance(computed_result.data, np.ndarray) + + # expected weighted average computed by hand + expected_weighted_mean = (uxds['psi'].values * uxds.uxgrid.face_areas).sum() / uxds.uxgrid.face_areas.sum() + + # ensure values are within 3 decimal points of each other + nt.assert_almost_equal(computed_result.values, expected_weighted_mean, decimal=3) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index bfec83280..5c07a70a5 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -429,6 +429,77 @@ def integrate( return uxda + def weighted_mean(self, weights=None): + """Computes a weighted mean. + + This function calculates the weighted mean of a variable, + using the specified `weights`. If no weights are provided, it will automatically select + appropriate weights based on whether the variable is face-centered or edge-centered. If + the variable is neither face nor edge-centered a warning is raised, and an unweighted mean is computed instead. + + Parameters + ---------- + weights : np.ndarray or None, optional + The weights to use for the weighted mean calculation. If `None`, the function will + determine weights based on the variable's association: + - For face-centered variables: uses `self.uxgrid.face_areas.data` + - For edge-centered variables: uses `self.uxgrid.edge_node_distances.data` + If the variable is neither face-centered nor edge-centered, a warning is raised, and + an unweighted mean is computed instead. User-defined weights should match the shape + of the data variable's last dimension. + + Returns + ------- + UxDataArray + A new `UxDataArray` object representing the weighted mean of the input variable. The + result is attached to the same `uxgrid` attribute as the original variable. + + Example + ------- + >>> weighted_mean = uxds["t2m"].weighted_mean() + + + Raises + ------ + AssertionError + If user-defined `weights` are provided and the shape of `weights` does not match + the shape of the data variable's last dimension. + + Warnings + -------- + UserWarning + Raised when attempting to compute a weighted mean on a variable without associated + weights. An unweighted mean will be computed in this case. + + Notes + ----- + - The weighted mean is computed along the last dimension of the data variable, which is + assumed to be the geometry dimension (e.g., faces, edges, or nodes). + """ + if weights is None: + if self._face_centered(): + weights = self.uxgrid.face_areas.data + elif self._edge_centered(): + weights = self.uxgrid.edge_node_distances.data + else: + warnings.warn( + "Attempting to perform a weighted mean calculation on a variable that does not have" + "associated weights. Weighted mean is only supported for face or edge centered " + "variables. Performing an unweighted mean." + ) + else: + # user-defined weights + assert weights.shape[-1] == self.shape[-1] + + # compute the total weight + total_weight = weights.sum() + + # compute the weighted mean, with an assumption on the index of dimension (last one is geometry) + weighted_mean = (self * weights).sum(axis=-1) / total_weight + + # create a UxDataArray and return it + return UxDataArray(weighted_mean, uxgrid=self.uxgrid) + def topological_mean( self, destination: Literal["node", "edge", "face"],