From 10bc54ff482827f76574bfa1a7cf568b5804c18c Mon Sep 17 00:00:00 2001 From: Andi Gerken <andi.gerken@gmail.com> Date: Fri, 21 Mar 2025 16:27:39 +0100 Subject: [PATCH] Added rendering zone sizes from attributes. --- src/robofish/io/file.py | 62 ++++++++++++++--------- tests/resources/valid_couzin_params.hdf5 | Bin 0 -> 33648 bytes 2 files changed, 37 insertions(+), 25 deletions(-) create mode 100644 tests/resources/valid_couzin_params.hdf5 diff --git a/src/robofish/io/file.py b/src/robofish/io/file.py index ca73444..c45bb74 100644 --- a/src/robofish/io/file.py +++ b/src/robofish/io/file.py @@ -29,7 +29,7 @@ import warnings from pathlib import Path from subprocess import run from textwrap import wrap -from typing import Iterable, List, Optional, Tuple, Union +from typing import Dict, Iterable, List, Optional, Tuple, Union import deprecation import h5py @@ -1097,23 +1097,25 @@ class File(h5py.File): if categories[i] == "organism" ] - zone_sizes = [ - ( - get_zone_sizes(self.attrs.get("guppy_model_rollout", "")) - if render_zones - else {} - ) - for _ in range(n_fish) - ] - zones = [ - [ - plt.Circle( - (0, 0), zone_size, color=fish_colors[i], alpha=0.2, fill=False - ) - for zone_size in zone_sizes_fish.values() - ] - for i, zone_sizes_fish in enumerate(zone_sizes) - ] + zones = [] + if render_zones: + for ei, e in enumerate(self.entities): + zone_sizes_str = get_zone_sizes_from_model_str(self.attrs.get("guppy_model_rollout", "")) + zone_sizes_attrs = get_zone_sizes_from_attrs(e) + + # Check that there are no zone sizes in the model and in the attributes + assert zone_sizes_str == {} or zone_sizes_attrs == {}, "There are zone sizes in the model and in the attributes. Please use only one (preferrably the attributes)." + zone_sizes = zone_sizes_attrs if zone_sizes_attrs != {} else zone_sizes_str + + fov = zone_sizes.get("fov", np.pi*2) + fov = np.rad2deg(fov) + zone_sizes.pop("fov", None) + + entity_zones = [] + for zone_size in zone_sizes.values(): + entity_zones.append(matplotlib.patches.Arc((0,0), zone_size, zone_size, angle=0, theta1=-fov/2, theta2=fov/2, color=fish_colors[ei], alpha=0.3, fill=False)) + zones.append(entity_zones) + zones_flat = [] for zones_fish in zones: for zone in zones_fish: @@ -1388,12 +1390,12 @@ class File(h5py.File): poses_trails = entity_poses[:, max(0, file_frame - trail) : file_frame] for i_entity in range(n_entities): - if categories[i_entity] == "organism": - for zone in zones[i_entity]: - zone.center = ( - this_pose[i_entity, 0], - this_pose[i_entity, 1], - ) + for zone in zones[i_entity]: + zone.center = ( + this_pose[i_entity, 0], + this_pose[i_entity, 1], + ) + zone.angle = this_pose[i_entity, 2] * 180 / np.pi if render_swarm_center: swarm_center[0].set_offsets(swarm_center_position[file_frame]) @@ -1454,7 +1456,16 @@ class File(h5py.File): plt.show() -def get_zone_sizes(model: str): +def get_zone_sizes_from_attrs(e: robofish.io.Entity) -> Dict[str, float]: + if "model" not in e: + return {} + else: + possible_params = ["fov", "zor", "zoa", "zoo", "preferred_d", "neighbor_radius"] + model_params = e["model"]["parameters"].attrs + return {p: model_params[p] for p in possible_params if p in model_params} + + +def get_zone_sizes_from_model_str(model: str) -> Dict[str, float]: zone_sizes = {} for zone in [ "zor", @@ -1462,6 +1473,7 @@ def get_zone_sizes(model: str): "zoo", "preferred_d", "neighbor_radius", + "fov" ]: match = re.search(r"{}=(\d+(?:\.\d+)?)".format(zone), model) if match: diff --git a/tests/resources/valid_couzin_params.hdf5 b/tests/resources/valid_couzin_params.hdf5 new file mode 100644 index 0000000000000000000000000000000000000000..e8c7ef32b4b929f0db806dadbfd266ab1eaca66b GIT binary patch literal 33648 zcmeD5aB<`1lHy_j0S*oZ76t(j3y%LoK|wP_2+I8r;W02IKpBisx&unDV1h6h8Mqig zauN_Og8<Zg1!jnV21t^DfgvQw)s=yPkpX5tjE1OUU|@h60Hxr<ql}Re0v@i80U)17 zfCvT#1`Q~E0-DaCT!z%VlFX9K)M6OFI5D>%Co?Y{CIC%t3<fX-1ZY@#eGvswm;s>} zco+g0SQtPlmXAS%L4u(?zbGdqzBscgH9k3)fq_9#K^`p6z@We&&cPsF57o!W2yp?- z*)W=c;e-<i^DuaT^)WK=Gw?7-Fr?)d<tCQIm!%dJXXfWIFbHrkfaO8vf{esv9>f*~ z284Os3?d91F!M@_azMHn7$g}P!8sC?>mhnLz!DHP0|PkFAwuAM3=0nj1+X#(Sa>ip zfFL}ZG6;fAWnd@(1rh@T!vr>ndP4c!K^dYBY9)j-N{*NiK+or}@_qv}{lEgEC_gE` zgn@w}EweZy-Vl@zAu<jSie!?Jfd!ntVfMk&`2l%|Q4J7^0o{F|d;^jSU<Z*1Y-FmC zQ3Wm^AYvX6ifj^>`#hlTtALu!$iRlx9~~+X1u%a=cnl2g&i+2&ybkpUn9Tsq_Y4dS z86_nJ#a8<I>6s;ZnYjgeX{EYJsYN-Nd3q_S`bnAj;CRSPF3w3z(g(3Y@u6Q*l$e|y z4=K9zb@hu<b5awFQ;YQt^$cL;9V^rUuyBGl2N<-W;vpHC#R{3l3W*BEnYjfysS2qT zppraQp**uBLm@XYB~_ucKp`cuBvB8n16s^5z{~}?1(ZHue3(4QeIO06<Rk^+FfcIq zRw@+bmn4>C=A|nX7bK=<q?VLqD!^1TFu<Z0rY||av??=?fdN)dp^I}u%|#c7G_)8P z(8alNsON!-!|cf~N>9woEY9VHildvu2Ng$G&kq$xR}YPLxIbX#3*t}@6=t{qa{x?W zlpYO%(GVC7fzc2c4S~@R7!85Z5Eu=C(GVC7fzc2c)FBY;=kLPEzy$CAc|beLuyFwN z@jKY~2nV!73{&SJ3Gt=@R2;$^IV7QRO$IqgSit(_5Z)*`QbGWdJ{WM%Z;h0I81*OC z5I`SShLsa5WKosF%3JWf45}1aJV^P^2rvIvK+A8KJsHvvV>{5?0TCS~VIcsU??MkR z9tH>Sd>tbLFJ%6$D7COOH7~g`9y~pOd!7(xFvAH)2op9^3*o^?n7!!rFT!39u)Qg% zX^EvdB?RaDz~kl6#0MFxht>En8&PS*Jg6Ypp45th)S}GX)Vz}T;?$DT0+?z-^Niqe zf2h5Xxd12!Y5=7C#JxU1LJgu2RTo5TlpOvc04qN?K)ni^m$q<)3P7oX{Nl`#%=|pq zdVu_*%+x&4ybn|eR_+nIJ^)gFGC^m%85jzn8e!(lfbw@3K^y=}pBK=|*%i?F^b<-D zbt8?0m1F4PfLOo42`*of6H8Ll^NT8B;X|lgK`zfgT^rcE?hR;sz}Az%=1swqr{H;8 z$O;oiCP+-CKrMuni3|+rd}am-uvT2_K3Ewz!17E?5Cb`&3!0c26VxF_!&XE=ctX$` z8D#suUiSnB_(}_iJgrGcx@KX3&0A}@;I|PHS_}*l&iK{AYLf(g{OVx&ghLg-y1_52 z-~@7G=UREEhFB$sifwXE7daO>oRyY!x@r5#LGQkllh_^w$5%`ePW_wR9Y1%9IO*Qa zb-Z6L;B>a5)A9ZVZYSHRvmN)(WOMRiT<duMI)f9-?!AsNN?#o{H=cF$p7Gq#?eZPR zI{BN92L)d{J~(#L@pbGEN9TDv9XUdoocy<~a6I^h-Kj`nx?>0juhW@#EspJ7f=;#> zDUO@UM4j^4?HtdXkZ{`fNX#+4R@#Zv^`677iLy>*_j?@{Fv~kln)lg$_C<N8HRVeC zZ~T&Xa(L)<K&M{bX_w}<17h)VPHn=>2M^4aaT4m*I{3>>%IRWI;6WQ>ai^0`<p&o{ z7Iyj|*>~_oG`|xA!-9k96S$m|95)^mFk^M<IeqZp`6d4x4f`%0WW4^-@p8q3gD?1> zI11FfJGedes$)gx?}O9MA9Xx^lKD_~=~hPucg{m42}>L|oADiby>OD_*~vnOtbH0B z&qs<Kn))Ql@oJ0YApw6A$Jv51ht~e*ajgF<d+3?%DTfX>`9tr*iyabQ%O5HzT4WEU z_d~@GK-C?9>N^NE=OEPFgHU@8LhU^Wb;m)dyADF#c@XOEgHV4Qg!=0s)Sm~T{yqc^ zheObCIRp)-L(p(L1P#YS&~QBj4d+A9a6bf%heObKIRuTTL(q6T1dYc-(0Dxrjpswq zcs~S92Zx~P;t({Q9D=5sgV1zz5Sp$ILetqnXu3NHO@{}e>GB{nogRdy+k?<_d=Q$h z4?@%V0cg5E0L=#n7#I|s>S_&EFfcGQFfcUOr#)%h%U4rxU$T&K-+>Rc_9CYp_pOVs zwy&7hu&?HHx&6&U^Y`5`FShUY+PANDWv;yl-<^F%YU%dJ-~HG(V@te!K|0s|l$db) z3n#?)8>;!)@9tIEUn}im&$`lJ{~|X_dlNpJ{iiqS+IMbr+y5+E(Vn*>VE>#t5&PNg zk^3Efu-U&|mbky=>_@vLcQf|8Z@gj`ppw7;)Xr^ol_e$nKeNuTTY92m|Ea7(yBL|8 z{Z}p++8O86?N|Q)$oA;g`u&biooo+$uiqc`vvJP@!}|SaeHZQhkW#z<_b$PGXC_tc zzrG`IU(CL;{f_w^`zF6G++R|<d|#$~&VKguNB7-~Pu;&c=JCGkn`8GgFfi=*Q4Za| z^(^mxp6%ZIZ8@d)pQ~`(|7D5#euZMQ{dXoC?@!pMwSW0ZyZx)}<@O)9^4xz-OK^Y1 zouK`(4b1z`E{Wd1dd{nTJGzqhA8|XoPqZy-zl-epeVh6V_B+^5*cZC6bidfl#C=D$ zR_>QOt+?;ZiJJZU)*RYvf4y$M+E>QCjQ8t7v@M9<x8{&7NL=o;q8&(`*vv#bkUj_d z33ecJTx8eVfy_PPcGeDL&+0j^>_GO$HZa?R+;L4y&>rNj)%J4sAa^Eg)UpS;TcOy@ z9^{X66^`~GfAMVhwg>ssM>*6U<nQa7W9>oVa5FyD9uzK_@;UaPaGLzO&>j?SG5gBw zLE(62Qk6X@TtB4L+JnORfnmKpDBKTxueS%q!_lqv_Mmt%&Z)Bp#Z!z-jXfydmY%4v z2gPG$Nr^ouUISF}?LqOp<ZgyNDBj;LOSA{2gW2to_MmjZ+Yw+7N++Eg-Rwc>#)Qwt z9+Zw)R~p!Z($(%>6?;%RyKq9>9+d71(z)zG>G1fwA9kR0DZ+Qh4wO#2z4qCG((TPd z^X)+CxMEs^9VlIkoOZMWrSp=7jCP=OpZ28D7L*T+40}~U`2tjTjN+ji0t45sY|w&) z+TdzeGBR+2+rhAQAgn!ZU<`3BY<(i3{{8_Ih<aE%5yBfKM@k5w_xE7~O&Q?UHfY_u z1$5kvn}LUcgP|b5IJG!FBe57dFaTOd4@o<a6~mx@7AWjQK*T>#cOL(E+YGdRJuDq3 zfNC!W28I_Pa~MGGfQ|n_#^H(!Qj<a36A1NN!Rt~X4ur221sOST<w=4$$bJTf!Bw6> z`l+n&erkjx#Fe;6SU=SR679H@!F$^;iXcj0=|ccIaAZ&l5r>sy9A*$cE)p7*3>%>3 z94>VOB|dQJs=*54=fRb(Aa0@t>2ZjG0fb!+F))Cz!yyI+5VkqQzyQJ)hZq<@*yIoc z0|*-&VqgGaokI)^AgpnSfdPb74lyu*u)-k*1`w7x#J~W;5{DQVKv?7u0|N*P9AaPq zVV*+_3?R&Lh=BowSq?EUfH1=$1_luRagc!lgg+c)U;yD42N@VZ_`yL21`xh+kbwb& zFC1iG0O1n{85ltLz(EEE5Z-Z+fdPa!9Asbs;S~oN7(jTzK~Nu;fnml$2%d0|fdPa& z4l*!+aKk|c1`sYd$iM)?2?rS%Ksewa0|N*<9AsbsVS|GV3?Qs<kbwb&1r9PWfH1>B z1_ls*aDagUgby5GU;yC-2N)PYxZwZ;0|+}DU|;}Yh64-?Al$H@fdPbJaoS)HZC`^* z0tW^L2PkfEU|;~@1r7`hAbh}qfdPacI504PFaxNBab#c+fMNwl1_lr|aAaTrVFyPB z1`rN#WMBZ{1V;u25H4_JU;yC;M+OEE?r>ya0O1La3=AMV!;ygjgcmq6Fo5t1M+OEE z-r&f<0Kz*Q85ltLfFlC~2!rN@K=^_q0|N-(aAaTr;RlWk3?Tf%k%0k(KR7Zlfbb7T z1_lsj0M(FA3=Aw#%;Chq0Kz;@3=AMF;KaZH!Xi!#3?MAw#J~W;GENK(APgGM0$~*= z1_lt;aAIHpVI3z11`sxIVqgGa6DI}+5Vmk)U;tqoCk6%(c5q@~0AUv=1_lrY#RCJw zy1IJ%z!?p;*1v1*qr(~Objz#lA3w0S6Z~3kKVe0kU13GBeU<q<yQ{2u_Lr*m*zI4G zZolvUO}j;*3HI45zS%ucjj;cu&1t_^+TUKHQp`Tu!PVYEPuYIMTr2x(JAE()u|Z-W zIglEVT96))UXU3evp{Bo%mvI}G<$Zn7uKyCrK3FJ1A8$oUbxf$g4^mNOmp!5_t zqhT+IA05uH4<!Hif&D&^`UxxQ_JQ<Qna|q?GVfB=o_!$m_x-=Q4`g5Vif{Ws_J7jm z+z)b(M5Wk%kozq3l=p+&TWzPmALRZGbFKD+{1ffqx*z11z0&^sL4J9l8nGYbmqnoo z`$2x$zbJh_$S+q}^Y(-MQdm*EALJLoujTtee$g$j-VgGN_3zsKAiu1utKSdu3n-7I zr|(<}<?n^cgZuze5Ap*@KgbUt^FV$8nGf;<$UcxCK=$v0x(DP3ko!P>0J#_B2ax;0 z7~~hI|G;WMeu4TQYzD|L(C`7<0rCqp{J?Gi`2`xjAU}Y@8{`L2`0qb;G8hyeAigcg zFCagH`~vbb$nT&q0O<#X2`FqpVFWTC6lS2X1BD^Teo&Z#`~nJNko!Pk4hnlv9Dv*p ziW5-WfZ_<`7f_sm;tmvtAisd(6co3hI0pFz6z8D02c-d!UqERBlr}(V1mqV`ngOL9 zP#OZIC0J=SN)OEtK+MBSfcqh!9cA&wC5d?{iA5>#IjMQ+B^jU{!?^k*u>E2@4CM7^ z&OtZQgXYmdAqeV!gZ2R@lQ@qCZoh&F22kgS0kQK9WGZO?C1l?rxO0zZ{tVWRhRvVB z_=x#4d+5xU1A_@SgUNr$VgLr*^JgB=DJtCU$IsF(N%WpSBh;Qwh6Eb}6O=~hlOG-# z-Vh@VU|apbT;%!0fd~(<BN!M!{cybV@5zZd$)!1oC8;U#$W!z2MTw9UgCw7r44%l3 zFD^(;O(~8qDJ{w?X2{J?Nd;|fhovh};|^vwDKw$@gSZBi?jfe4^U06D8G#TJN7FMX z9GDq*gn(!UhLM?`Vd)rJ$U%Z)lpGC#!5ab#_KR)+je<7p7d-%?6Xq+N0MQHfi(UZH z4f@hoK(xbrg&Sa+U+oTvUa(*E0f=rWmv{uG^`)PH=!EC8&%pG2g%=>&Azb+tnC4e| z1E!B_yaUk-_KSW1(GAk#AHj6F#3wL)S@JWO)|dVQrn_ang6M?jvfse8v;223JzwDm zh<5m`_!CTrEB^x1n^k^;X@0doAX*_?{V$k4uJI2{D{K7+(GB}W89>!dL+&&PEiKLn z=AWDnq01$hz<kx25c;wtGnn5v3qtElvw-<`XG7?28CDQK!E!EyelE)f=1-jmp`GQ~ z!TgW&A@qC&4lv($A%y;|$O-1JTm+%RmASxt)+G>nvkEtupSTo4^Q-ZI`FodvX@_ie zUJ&0was`AwuE7W9m#&1+%3Azj{>4>bdcl5C0T6$ISQv!Pod%(gT!YZk;(}oLqHqX( zayo=oxB;dc$|ZzA@(t%AAhhaC2wig%LSL2?2Fq(jLFmR=5c<X~2(2$I0+w%!hR}Cs zLuivb5V~7N6fFNJ20~lTh0qi3Lg?qRVqkf@I0!v;9)x~#4@@UG%Zr2L6K2Ij=#TRu zw8sMoJzqfrEdMP5Li;X+&`TbI=?3Ya;PeU09}0pG!Ra&snqLyWZ3Cy*15m!h=Spz8 zh4CLaX+qL3RQ`eS5ip+t#@F-$r)!1-(EK-d7YbiNRur7xVd@=xBO!dKdWV;75I)qr z1Bd29_)vK`A1V)550ytZ4{je+z5(W*eGvCOKyxqLKT!2B|HAwO3lA6{7G4Zc^I+i# zlZS;jjL*Qp01khc`i7FXZ$bGMCf~4T&ubL^z5JKR`~|U`&%u1CeGSvP!Q~TNKkv_n zV0ozf8Fs9=2j)Z7D?~fo1oNTt4b}TDqws?Q&!F%HIS(T98M5{*L+69mjKJ*|<UEMN z4+=bk!mr+U8JQ1q4>BL*UNC<TBz_ure?A2BVeXyQ{REj03J+v{!@c~MDEu{hUZe0! z-o8cVgW?034~j2jJ}5qs`Jnhl=7Z7$G9Q#)Q2542zM$|my}lyzLFp09hx!+kUcr2* ze?jRPh5xec8!{i1AMo)(`2$%V-8^*r;O>Fi56U0N`a$^vh3}+^;y?88fQQ!(NPMD) zCwh2e#0Oe@f$|5k{pj(Hksi>}3wnBjr#Gm7(bFR+f1ucho}SUtJ1Bo3s|WFs^GgF- zeriC^Z)o`u<Q}koX!wEL3+6-J3-S*NAFVt9g$If}T6qErPZW8y@+bk?4mt2a719oQ zrU`C`EV!);X@^`i1h+#L95aQqL-tsL+aU*T$U)j4mluF(hjU8c_6NhsC1Co%VKqqm zW8Vre?XXh|-2PzLvIa~aSf>YWe<ZBj0Hzg|7=ha#59VzF(+g&rLE0Zvc7SPxek*YM z<3Q&gFx}8%2X22j)E@xT3RO<v_Q!*gBVc+#z8kpx;gEF#Of#f<f!iMk;?IEThA2Ns z`y=!Mm}UqF0=GX7cwGV04X$C}_J@MQ4KV${Dhks6Fuen&9VGO@?TrS}0x-=Wyb(-0 z2z&z53wVve?F|O*5-^>>u@y`oVEqQB6`0K-?G1(sF#X`~4lwQT`zM%Q@WTq+-eCAv z4W=DF?*Y>XKKupK4Dak9?Ty#<VEVv|17KR=86%{f@z@DWKX}j#rXB7b0n-a^vx3_h z3^(1tbi%cEFn!?i2{5g2krUj`XgKEurXQT{2Gb5F&w%L#$9Tc*j0cDPz_i1`elWdY z-vuzuut!i5)IM=wj$Z*zcQ^C)fanB?@M9qQ!QrxVAezC#8(c3P;Elcx;xBldbsI!C z?5w&Eq7S$q1=n*8MrR*`_zkjGpMq(QyWsjU;qw!4z3Fi84Y)o{IQa!!&nj&B^B$yb zfj<+t-J{^j25#>(SaE^dISdAT;Py;|st}~zA}t1Pe<TP<f!pT~SmnU&b6CD$*x(IL zpA2F4;QodKj2{8vA6Snn4_6PBe*lw*@)OYX!_9-5kFI_Jx_&hN12pr{?MLUMn-8}i z>OMF>1meF0H2WIR>}x=iNB1AP{qXRDnvX7z?moDD48%QfKGePFe0X?4&4>F3sveyW zcMnt^-8^)8bpI_t3lF$@sQu{Xq5BUWzEJ-ifSDHn@vj3cd=nsixPGWSoDX%s0!%-Y zkM3Ud_=ATRR6V+WboZn4(d|bMU%3CE?uFYA^$%P<R32_WR309lP<eFw;o$>yKRo=P zeDw4QHxDWgw-3rkmxsF-DvutX=<3nkhb|BIFVsA^{ZRiAiVyVghS=xEy%DYaNC=0d z7l=Fq!vQECt{=*W$vZ&Z1NSdf9u|KMP(Dl^#)ru>K=s4@2Xzm6c){(5s)yOn05u<- z4=>*$A^w5e2Q?oazEF8|^WotORgdl-xPEl^K>2X<p!(tFLFM7@1@U3~%h1oKM$C)x zfcKX{&+*MjEJy{9un}6P2Ri5RMG{ykZ0ZZb86`(UU^E0qLtr!nMnhmU1O_|=aGi4v zs=q;wD@si+Nz6;nfz<jS5eU8jJ2)A_9VJIYU^E0qLtr!nMnhmU1V%$(Gz3ONU^E0q zLtr!nMnhmU1O{;kFfuTM$LC@DwPE}9S3vi1!|t&pbne6n=tW!5tuTYw$D>Z8MF_yo zoq(Mi2un`^&~yZ|vmmi3F*mg&wFr7%0BjvUXaNRH4JkBi{tEs40nq+3kU5MDtl;^p zyu@7a5Eh|%EAZwZXd)H_@7H3-nm+_Um)JozP=l<2-qSjIjx{(iMhT1%K+hl0B*stx zEjM8rVEKVo=U8);L2QEAV*up`pt%F0YLtY90IXhAfZq2Eiys~a$bmPYdniHY=~U(C zgS8Oew~DSW5wAY*`8(iqbI|oc4r0M&ANZVCu*uXVVEb=j?xK}F)OGc!X`>-98Umvs iFd71*Aut*Ol!O4R--X_OL~jqJ<wF{nxY{SMa~=Q<&+Y*L literal 0 HcmV?d00001 -- GitLab