From 7be2e410fc2347c31e0d42b65e6f2d2b8e48d0a3 Mon Sep 17 00:00:00 2001 From: Zhiyi Wu Date: Wed, 26 Aug 2020 21:40:13 +0100 Subject: [PATCH] mbar overlap matrix plot (#107) * partially addresses #73 * add MBAR overlap matrix to the MBAR estimator * add new visualisation module * add plotting function for overlap matrix (from alchemical-analysis, code used under MIT License) * add matplotlib as dependency * add tests * add docs * update CHANGES * update AUTHORS Co-authored-by: Oliver Beckstein --- .gitignore | 2 + .travis.yml | 4 + AUTHORS | 4 +- CHANGES | 3 +- docs/api_principles.rst | 2 +- docs/conf.py | 6 +- docs/images/O_MBAR.png | Bin 0 -> 9960 bytes docs/index.rst | 1 + docs/visualisation.rst | 39 +++++++ ...visualisation.plot_mbar_overlap_matrix.rst | 18 ++++ setup.py | 2 +- src/alchemlyb/estimators/mbar_.py | 19 ++++ src/alchemlyb/tests/test_visualisation.py | 27 +++++ src/alchemlyb/visualisation/__init__.py | 1 + src/alchemlyb/visualisation/mbar_matrix.py | 95 ++++++++++++++++++ 15 files changed, 217 insertions(+), 6 deletions(-) create mode 100644 docs/images/O_MBAR.png create mode 100644 docs/visualisation.rst create mode 100644 docs/visualisation/alchemlyb.visualisation.plot_mbar_overlap_matrix.rst create mode 100644 src/alchemlyb/tests/test_visualisation.py create mode 100644 src/alchemlyb/visualisation/__init__.py create mode 100644 src/alchemlyb/visualisation/mbar_matrix.py diff --git a/.gitignore b/.gitignore index 2a45ed72..f71e81e2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,4 @@ *.pyc .vscode +*.DS_Store +build diff --git a/.travis.yml b/.travis.yml index 3d9c07ef..d4f03437 100644 --- a/.travis.yml +++ b/.travis.yml @@ -2,6 +2,10 @@ language: python sudo: required dist: xenial +env: + global: + - MPLBACKEND=agg + python: - "2.7" - "3.5" diff --git a/AUTHORS b/AUTHORS index 672262b7..c48f9fa2 100644 --- a/AUTHORS +++ b/AUTHORS @@ -33,4 +33,6 @@ Chronological list of authors 2019 - Victoria Lim (@vlim) - Hyungro Lee (@lee212) - - Mohammad S. Barhaghi (@msoroush) \ No newline at end of file + - Mohammad S. Barhaghi (@msoroush) +2020 + - Zhiyi Wu (@xiki-tempula) \ No newline at end of file diff --git a/CHANGES b/CHANGES index 98638667..b95aa851 100644 --- a/CHANGES +++ b/CHANGES @@ -15,11 +15,12 @@ The rules for this file: ------------------------------------------------------------------------------ -??/??/2020 wehs7661, dotsdl +??/??/2020 wehs7661, dotsdl, xiki-tempula * 0.?.? Enhancements + - Allow the overlap matrix of the MBAR estimator to be plotted. (PR #107) Deprecations diff --git a/docs/api_principles.rst b/docs/api_principles.rst index e644fb37..e0c6d73c 100644 --- a/docs/api_principles.rst +++ b/docs/api_principles.rst @@ -1,5 +1,5 @@ API principles -============ +============== The following is an overview over the guiding principles and ideas that underpin the API of alchemlyb. diff --git a/docs/conf.py b/docs/conf.py index 517695f4..28d24289 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -50,8 +50,10 @@ # General information about the project. project = u'alchemlyb' -author = u'David Dotson and contributors' -copyright = u'2017-2019, ' + author +author = u'''David Dotson, Ian Kenney, Oliver Beckstein, Shuai Liu, +Travis Jensen, Bryce Allen, Dominik Wille, Victoria Lim, Hyungro Lee, +Mohammad S. Barhaghi, Zhiyi Wu''' +copyright = u'2017-2020, ' + author # The version info for the project you're documenting, acts as replacement for diff --git a/docs/images/O_MBAR.png b/docs/images/O_MBAR.png new file mode 100644 index 0000000000000000000000000000000000000000..65effe53df2696df1c618d0fcda6e4877756c61f GIT binary patch literal 9960 zcma)ibyU>-w)fB?rJ}S9JcQCINP{SiC=DVZU4nF{iU^2+lt?2nG>Fn6B`w_{A>Gn- z_dMsWv+jA{bMCr-80gGzzOnbG_5?qAC`)vi>M{a>Ad;86uL^%Qz_0pCxbS-oEuIkk zA>{Nx+eyvN%*oZr!4#ov!O_z0IUgq))U?JG@Q>blJPjn@7EI< zU*&1NA$a?${n2F{%HIK&Ll#u&vnf)`^Zd;t+QOfI^QLK9h0gN6jShZN!FD(EWryMW zfVpdbvA(V-(&wDmJqw%`Gibx9pxG#!J`+Wf(2{HjE&AvqIS-tNdKS+jO>|<49lDdm z6kFfjd{rC0GOn$3#B(pJAK#m8h@AO~w%B)U%v{+xNyEKMq?ht&jWGQO8I;cZu4#$kG+*wrKNl-qTASwuZb9L-DrTTE&XX9=b}Ryn48 zFT5Wni3S%Ia`i}Dhz=JT2!{B`%j4J9)_NT;XPNb$+{wwwX=-YE9~09uRpZ&ax<^RC z?H3t&t#rxL-JQpEL+|0IyA9iOt?v7y=D}~?%(NnX`9sbpX+&H?Z?I|pOqV6;r+U9C zWoemXV{FV|H&tCB-W7LWQqqrxlr+7%0ELPa6BAnya{N6|k&pB_Ox5?sn4h0FiJCSr z@Yac0SRF3x@cZ=fqt#XPPi9>o+b&KQCAv@s>gtTHOFyNis$2_cXlU3~;(0%J+YYRL zOerb3CMzpjQc|+guauaRojubOO#1TGE0zu$8=Ff61anKh>1OL4bKlp;zhPLwaz@K? zQK*sf7Ys~HOn){e$MeQ`cz6h@__JGFQwuCs`uz{4Ji}*apOq%06cxQkp;#v-Cal<9 z9UU1&M8@;H-D$Wf2$d3rqielS`tqL$mM)u_ne|n<+O>Cdyl`;%OGqf5^J!(AvhQ3= z*e6lKYAJD+1pDL;Gm>_>b1?AEVFv%&R&98yY3%86R< z1my&QtCW=TR#qsT(D$*itshv_Myg!7h>3~!tU5CK=>r=iot*fKi;JuG$E;5e>MkA> z4iD!tI5!Lq-gxrl$>qzJt=PG4->#pWoU}P}qj?h+W(w6yBkFdAi;Jt&%DPNaNrho5GS_f|#E{`6b$Q}I&~K@qU)lwgO+NNZ>eGJG?Q#E;>8 zs$gV9UuxEs=6SHOwC7Qp^xC3K*JFRJ?fmQ{+7vg?|Lo)-;q23Z>W>m78yohvwl3yX`p-|Ei7 zGBOyc1fMGf5#KEO_U#rMTgb*_)o__ml>2J_w>P1oP4HGfKg8_Z919DJ(M0(RvtBMv z&Vm=#hqKKgCMG67tKA)|Yig+VB$%0B6HyBkIj?GdN=&S9S*JoMf8;-rvv6&VB>0dLJ429X^LzkGg(z^eAaHxh0IcKkY$Cr2+Zio=YFx(?YC)Pelu_ zgoMOVq3@Y>qMh}SQ;_Nl3<63yIbEKJP#FJJhh#&{OUCL03rOlOjkfBcXuU7A** zeosnD+Lb2ZbLv(B7mR-Oldm=`L9bu8d+x90IErMhOp~bDMGP&=6&be@6Vu$ShwHIo zhX${D!1>PD*f=;W>`oP@q@3KP5}YjaC#uwj3+>TlgLOza0_K{+hs+eBZrp;#wQ2vh96*v);MlO5iMF&^iBMpsL`S~0L1qEiqTT?Hg5!>hekWn{eqE90Z1I!bl=j-jg^RRDDL{)V%%r*W&YksNogb#!%g z*$mRt;y`njTF1o1VA9E7>-u_Z zs07iHLyJ>L_*9QNYZd6p>*!Fzy&0Vy@0ktPIj@qzo|WRb{mD3VVx-K~@*{8BaZAjwcsOP?n%rj$S>4y)Ap}EyicXFkmwB+TMxpK%4)5*_8$=}}V>Z%2+72OgvoNo$Xenv7Oz;&wVBFYo~L_3&Yu_ z;81G8cK}R9UPlfzH*a?FQf#RB9+NKiBo`L5v$Kc3m}}_;5UX@qZx1HDb-0usu0F`Y z!NJ7C6ERw9&Sy6zI`-98QBo3t$k#2K^*vwo)h;r`K>D8YK$||B{}F*&5fK*N8M7|j z?iSe;wNC&{7Znv%AFN5Jt`>!>Y9crAYVgOxD}OPEOD*^uJ0&kKbZN-~O-lP~qmA8( zBHITCS{f39g0b*{yX)Vk9d_pB3k>SIR|a!2mWJ}F@rY?=rZ3I~r@aquzIyd4TpiiQ zNUF_?Oi_&GB%~Gy1*#j;-dJ7j_|+K1!pVs{Gj3OBx;9dbvNEXkiigic40?V%Jw2sz zM|}PIwOzlpw>RR+x92%1J9e?d6xPm)iV7Mr4|V6=$u3w8J`vGp(~cOR+3~S4l=yYy)uH^CK|yoBf4^pwN1twC-8ohEd1shx#@n}-fenzBqLK2< zZrb90#mr{XDZLhQa&mOS!pBLdMj17#lRqoDg-`)~hJk&6PTRY?YOI@pfR0-;QsRd* zSU`Obpxq6;_cQ`6ro-FY709^sJE`reA7?%L;tbu-_Qh&Y!F_#Up*fmO3xO^q-TavB z?Ay2`w4Glp`#&VU*42Z359L2$E3B!hsqi`~**Wg}toUbjIDCJ7d~SC3`%8RME8AT` zXV;~(gK6{K-(3g!eu;pC0XPIB#m0D+mX@bba*}(o(zei$Z{EE#o38bqUs%BSnf8Fd z%iB9%$dQ@Ptdow3soAF4A7iM{KqXn+yHl0$P)c1>b7#_Zdiz%pvc&tuWnywNp2zsL z&dJRBdKcVJTX(liRzHhuPLNR~K6^CD6&f1!-Wr-i$c4NPxj8vG^{U+-#}k&?jPimI zdhqDc73g~S=*Q~n8CJ-(_4RmBci#B;cwn|iI_OJ>LnwD#Bu2oZzgQ0!NFB7hM#aas z!A^riLruoMSbm6!xkgEubIcN-m38ar=%{_QBbIA!c{wMK;PU0~lan{)Bj3DvL&v}% zt>Fz^Y1)c4=y`s!z*E7RqGiV!yM{#22LeH2D5csRM~lQki<0;PGkU`larB;A3p|FprEKo$!A8Gnwn}RwEi|E zgn^&`0{~*CQv7$I6+%fvL%MXCFicc9cGN8uGr*0Jk&&6-zca^fRtipgMQv@_>%0BJ znesk(`n|CcV`;cBdFKH3_y)8knt?PlZh*>I1av4UTleTys+bK55H}1vt&Q~NeRnmsw$AN5_4rvm`Z}|^TJLCb;sX0n zVopBtWmHtu=w3)jNQ5b+eB>o!Vu7mjL=o59Y;1EmSxLN};(;#$Y9FWs={MpU8g9kd~g_RZ#b=`7&(9VxpV}Ez3vurad>s-oAbNo>6`e|895ydj}{( zKmm9`u})c7D&y>T?f+_h zW5GwT4c2w~f8NCaS=uwx5E$LwpQVy)2D}O?GV5WYRcH;U>NL=3&>*Nx-~Bb6uXfY# z9%e~LGFCb*e}`i(DlV4F>Myn$WZl`>0danyzOES(5P*p$9e9(yqoeV1-{%7xh7MkM z_98H_*LE{G>VO`kdQ;PV6P--0Lj9)ZW+QX+paw%aaq%P&6lehiCJEl4#i}83DN2dk zent|Pn8uif}dHw$LsN?WOALB|#Sh z2{rfjJ}ztbBw1+$NM_;SFkK%jJD43PW_F5dCOkZ?^90mdHWsy*qV~we}h&8{sD4{<6CM1O5rm)k@;r1MA z1za9WbI|(Q+SbmFs?4kvKb5e{+Q`gYGLY2rSlKF8zk!3Jqr9S`by4U`@00zjG&G^P zxonY)+kJ{$#n!`|DaUe@6cld~6K_T`_I|T7s6G0fpgVr596$4%J&|N)Z7mee74Ep! z{3#qAz~f^~X{rDb*NyiP5!-S&=3Ng=O__kwnFIthDPz+G;BxXKn$SRP^6c>N@cqY+ z3NK#pc+@0`dG6mH#V2V@!&*}M{ME}39`Uoz*B5lVte02y$}!mP`wt$xn5p-tGW{Ng zjg9^O)2F=3f{2|R2l0zj$9vM!93N%3tRo{Mm-@3v@}GSB`lDzK%(%neibh&m8rpks zbHD%c<%{n0K?0xD(;W<%cQiTbw>8xygoQum<*}=(scq^^7d#bNTU*4{&lG?!lSnloZH8D3Yb={mgc@E%v zAH?6@?fo-JC*hkwX@JPGd-vt;l#7r6EBFDbs_oIq8nYb3qtOK-U^|xKMJAj`ocJ&Z zE*`Ll5k%gwfp5)_fHJStix)4#M3CfM`iu+=flx?-KYoq&9U%&QGp&R$*j9g5rY*I1;B@SO^~<9|#f< zejw6RRaNh~y9ex6rKE8JIX1)m{NbDUryR9IJQ>hz?cs!23l>!0- z@$vDEpo0{crGSYjd7fcEI;;6~_r+h4jPxD6z$YFV!S|RKNHa1rCJr9a0Lym{xnBqQ z*_83|LNKSLtxbOK$Ea+>-JDF2(Wk2h7ra&Hu(i<|&uYZ(*kJ?SoowU|5wBJB)TXL@hKWk& zj-4ZH+`xusp7hxd-|%eUcRi(|$>zc7zjh1}27>J5Buw@%YHI4z1Ufdh+%ar3jk6mZ zPrlBBMC7*~3f)^Bev&}~?DQLM6U2-TBhv8r&+>J0^8WG{##8Q#H%nK2Pq$k?eE9Ga z{=uv>&dB8qRLo;Nz0k0*F7FUBenu)~@G*K61phAN`lq#OWhCv4=L^!W@%Xj2)=8=& z+E8TFXr*pKcF9LE?oSc>ojcKSaTI7Z2Mz9fwrW7hZHNFMc>Xi%9S9B2 z8{J^JJbDC5*h#^XLcxrDv5Z0;$8_VB+yBCZLW54g%r z#fvKS92|IQX=zdBSQL94@gzTiRRJvIK~e*zqvai3_4a_83d{H>ac`<8PsVY31^12q zEcJe|=*6tjSGji&^Oi~yBpb*unMOFat*xzkLvHvBE-|fWw)zh>EJzIXy|y0!MQ^Zp zfF3mhGuISK$yZcXhNgtI@o$Ag{_ZF}Qb>`gQyP$-P>TJKPXT3bREdMDLhBQ4TQ5L; zXyV(+#Q@M78ygeV?rJ&)dU~Azd{j>SX1SXA7QJb&`tX3W7x*kjOFpX#Ou8|aSNU4a z%+KEkv-jakT2)&|M`LR%UYJZEMFKER0|^#!8=PsCFtB|?BB8gZ}kTIInIsBw5C5}o=~3)F&|0Q;p= z=Ocz*h(@j^0|t#!9i6ti^F1ygIg6p2#^7%%z@Wmt zN#_oN%We#pIM-f1SuqV58PNtFY=C0|VnnH^?X_rA`=0EtE7+a_d4o4o0t2qmU+g6A z2ikbD(ivTVkP)(WaOqe7fZVqUlwY)IviCkEoK``^+TV1Mr&keK@U&{mkw{xZqZuMN zfM27fpQ)&oA#j!V@87?2^(uy>heu^elp6(eJZQ@)w>hH7s3<_$R1or+{OPQ&cz37rKHe-sRm3_wVDg33V>2xzJBenGQh&r+CUMt2iksTHiT>FiSORZ;Nj8H zFSrlS&6*%sB7_IurhO9Nvti{BHnRY7w1S<1cvM@P0?4N9t8E|69G;LU4^Qz!BO-KY zfyP6-sfpPrE08kz9;#uWnIGI5gfM+8-1oFGZ>-e3=lAzl1ewXi_}CY9bS=dA|0~}8 zYw~a0>xDTtk_xtcd(=ExCxZk|Xcn#jlC20YKgg3!Dc&;Wss@_wUQO}%ogxHM?lTH? z8SMv(0sHj}@lfXUpEEMFGe}x|{n$IM-?#w}k=Z!EryEe!(ewmU(fgdN>lSx6sk4)l z30QZsWoZzBIwdCfk&N-)ry`>IN0(wAuDJi4*nSN|3i^%`(_Q>fuvfREzIAnhPQHEn zHr?I3+Jlr9;q7re#^uhdG0nReRClbJkQaxtKmzASiyxw+f6dIKv`>UoqGMBtUpP5K zaWMb?(Z>HoeYEDG74v|RL;whi8qWhWu%=cA%5>ZWAoX{5cAkMrc3S$0p%VkVLUi@& zRkW$J+``N=@Ht`v$s!wL8W0c=cgH$L^+NpCtsn?LRkgJ*)kBP-o_}5I(W80@V$tCL z{P}bCrOBb{YrHBhVz|?(6I}IB!PtVSB zP^ejen7?x|UY9j(aK<{|)V;p}7ZCvR!s7Kc{$6;>>H0tH#Mbd@VSla`S6_cW1SHHN zB2PwLAi6`x{2EUs=;{4|vml(etJx!yl9Gb3>N*mMgmY;8B;*K9U%mC603C#Z3<~`S9T%61 zP@k5jCTsAWN8Ya;HYS8jB`ttYVWkj1l;HF7@qv^3J~fpDLNwyjC#6`rDQWJp_BjgtFSp6aXqbk86HjoF_Hoj3jAey`keGQ&r#hVkPc{t z!}F6(bbb^O5dm~OHt>5z86*ihR|M-zKmN_G6~-rdSbhJn20dbB$J7C>^wjmRb`Ss# zGxZqYG&gjRkQt*7XITW225RaK(EV*)T{x&h{DZYoX&8Y3n7_p!!fliu=XJlH1ELxD z2=QlaGzx}7#ctcrH1o8LtfI*vzL#Z@2pwnZ{w z8&kB7g_bU08lFd1}4_ zqX+mQz2%`4QOmoV9Uws@(F$n%D+jy`Jr0v}R*&wz0_yESBqJDd@V#@8Qv8Em@JX4A zd@pt(YSo3!(MH>=OUBa%Wt7~=DHu0e>RIlsCgys9k;blex z>@XD@%F{6$7K{nxf&d8or(jj>i*q|LU65Jx<$irJuxlKpVwi*B)ZE;B3M~V$pV}Me zn+&N3j(S1^tglR2&d$P$`e6EzHh*%xHhx1xgH$}Z;d2EA1xU1h{T(WKdm*&6w9sA} zuA`HeV|N(x4B4Z}YrYqg23|0YEVZI`l?-oQh775!8_V`KAud)J)txjT0W$oyOC%(V zYon!LbS7NCdr0qQNt^{k5rhz$;%a|;eiOX($k-U`E@|fjn1X!%{2B5ly8q3Mt=I;D z`2S1(8eEn}stBh4QOcS8%Gp`jR9IFIMP z4M=9>SZ14S0Eh;Ci2tu$yJj(z7oXEL5IixVXV#O{1s&A^jZoopCJZ(kHwCmdderKB zaWd_7FzI@oj0~lZWRZ|a?h!u;1^ruWI}SWh-T(&Be?>1~ZOsOdZKT?rPe1^Au1HG5 z`@h`<^TSVWJ(Z0?B#z5{OyC3?VBX@&C|{ZNFpG*}7Xv+)5I^0LhHxV5gM{F3tjP{JwKH;RjR5YH%Jbzi>tVl@ z8biQf=rxeRFziY=``t6?-F`KnK#9MJ=<)=gcD{RM?9GCeh%C;^2OrgcBp5856%I zzNr8c8Fb*l%K8>QGR|4u1UwD|d2E2?#^C{fdC{>N3QLG<)WP*97K<(PZLYtRY4&H(DLKpp$_3VjYJuh}RSafo5>j7Us( zbcBSEbFmhpSvWR!SZo|7sxU&auRUI21G%31u{J13Hcav@n^wVX_e+tCJeGZoKo8jU z;DjuO3lgoy%Af*(Ccyt7-X|sLGYrZq4|L;kN0*fFx@=B$5pY%4L;Q!(E`IhZGJ?mj z9>J+s{>WOyc}fG+O=Nt$yo}82ow?1;p4`G(1sNHzo@a0r?eKgM!C<1Uk&;r1d)-Bl zbLwUD$S=Z_C zAv*$OHZn2-kCl@~LX<$Vb^CPP&X?iNof`n~J^(Am+!PQrFacb?e2JZMOh`vhj}h3A zGLb0awDc(7HsOQHmi4`R_qJe)`JJBjhOxMRP*70wPB*M-6q*O*=hq(@A0d>FVb}SK zc~IM4HH=k7btCo()0I=rE50LCcs*+K$15&$|l zVw(uv6xbRF)?d&t5drlr5i!u`pm9rAFJOMz?uBrSon`0ZYBl5*rDeWyi9Pu4&1>3X b7X>> import pandas as pd + >>> from alchemtest.gmx import load_benzene + >>> from alchemlyb.parsing.gmx import extract_u_nk + >>> from alchemlyb.estimators import MBAR + + >>> bz = load_benzene().data + >>> u_nk_coul = pd.concat([extract_u_nk(xvg, T=300) for xvg in bz['Coulomb']]) + >>> mbar_coul = MBAR() + >>> mbar_coul.fit(u_nk_coul) + + >>> from alchemlyb.visualisation import plot_mbar_overlap_matrix + >>> ax = plot_mbar_overlap_matrix(mbar_coul.overlap_matrix) + >>> ax.figure.savefig('O_MBAR.pdf', bbox_inches='tight', pad_inches=0.0) + + +Will give a plot looks like this + +.. image:: images/O_MBAR.png + +.. [Klimovich2015] Klimovich, P.V., Shirts, M.R. & Mobley, D.L. Guidelines for + the analysis of free energy calculations. J Comput Aided Mol Des 29, 397–411 + (2015). https://doi.org/10.1007/s10822-015-9840-9 diff --git a/docs/visualisation/alchemlyb.visualisation.plot_mbar_overlap_matrix.rst b/docs/visualisation/alchemlyb.visualisation.plot_mbar_overlap_matrix.rst new file mode 100644 index 00000000..5323d1a9 --- /dev/null +++ b/docs/visualisation/alchemlyb.visualisation.plot_mbar_overlap_matrix.rst @@ -0,0 +1,18 @@ +.. _plot_overlap: + +Plot Overlap Matrix from MBAR +============================= + +The function :func:`~alchemlyb.visualisation.plot_mbar_overlap_matrix` allows +the user to plot the overlap matrix from +:attr:`~alchemlyb.estimators.MBAR.overlap_matrix`. The user can pass +:class:`matplotlib.axes.Axes` into the function to have the overlap maxtrix +drawn on a specific axes. The user could also specify a list of lambda states +to be skipped when labelling the states. + +Please check :ref:`How to plot MBAR overlap matrix ` for +usage. + +API Reference +------------- +.. autofunction:: alchemlyb.visualisation.plot_mbar_overlap_matrix \ No newline at end of file diff --git a/setup.py b/setup.py index 0e3e3531..eb689366 100755 --- a/setup.py +++ b/setup.py @@ -29,5 +29,5 @@ license='BSD', long_description=open('README.rst').read(), tests_require = ['pytest', 'alchemtest'], - install_requires=['numpy', 'pandas>=0.23.0', 'pymbar>=3.0.5', 'scipy', 'scikit-learn'] + install_requires=['numpy', 'pandas>=0.23.0', 'pymbar>=3.0.5', 'scipy', 'scikit-learn', 'matplotlib'] ) diff --git a/src/alchemlyb/estimators/mbar_.py b/src/alchemlyb/estimators/mbar_.py index 6d3b4e36..dd8127ae 100644 --- a/src/alchemlyb/estimators/mbar_.py +++ b/src/alchemlyb/estimators/mbar_.py @@ -45,6 +45,9 @@ class MBAR(BaseEstimator): states_ : list Lambda states for which free energy differences were obtained. + overlap_matrix: numpy.matrix + The overlap matrix. + """ def __init__(self, maximum_iterations=10000, relative_tolerance=1.0e-7, @@ -98,3 +101,19 @@ def fit(self, u_nk): def predict(self, u_ln): pass + + @property + def overlap_matrix(self): + r"""MBAR overlap matrix. + + The estimated state overlap matrix :math:`O_{ij}` is an estimate of the probability + of observing a sample from state :math:`i` in state :math:`j`. + + The :attr:`overlap_matrix` is computed on-the-fly. Assign it to a variable if + you plan to re-use it. + + See Also + --------- + pymbar.mbar.MBAR.computeOverlap + """ + return self._mbar.computeOverlap()['matrix'] diff --git a/src/alchemlyb/tests/test_visualisation.py b/src/alchemlyb/tests/test_visualisation.py new file mode 100644 index 00000000..d482caad --- /dev/null +++ b/src/alchemlyb/tests/test_visualisation.py @@ -0,0 +1,27 @@ +import matplotlib +import pandas as pd + +from alchemtest.gmx import load_benzene +from alchemlyb.parsing.gmx import extract_u_nk +from alchemlyb.estimators import MBAR +from alchemlyb.visualisation.mbar_matrix import plot_mbar_overlap_matrix + +def test_plot_mbar_omatrix(): + '''Just test if the plot runs''' + bz = load_benzene().data + u_nk_coul = pd.concat([extract_u_nk(xvg, T=300) for xvg in bz['Coulomb']]) + mbar_coul = MBAR() + mbar_coul.fit(u_nk_coul) + + assert isinstance(plot_mbar_overlap_matrix(mbar_coul.overlap_matrix), + matplotlib.axes.Axes) + assert isinstance(plot_mbar_overlap_matrix(mbar_coul.overlap_matrix, [1,]), + matplotlib.axes.Axes) + + # Bump up coverage + overlap_maxtrix = mbar_coul.overlap_matrix + overlap_maxtrix[0,0] = 0.0025 + overlap_maxtrix[-1, -1] = 0.9975 + assert isinstance(plot_mbar_overlap_matrix(overlap_maxtrix), + matplotlib.axes.Axes) + diff --git a/src/alchemlyb/visualisation/__init__.py b/src/alchemlyb/visualisation/__init__.py new file mode 100644 index 00000000..4fc73051 --- /dev/null +++ b/src/alchemlyb/visualisation/__init__.py @@ -0,0 +1 @@ +from .mbar_matrix import plot_mbar_overlap_matrix \ No newline at end of file diff --git a/src/alchemlyb/visualisation/mbar_matrix.py b/src/alchemlyb/visualisation/mbar_matrix.py new file mode 100644 index 00000000..682b8555 --- /dev/null +++ b/src/alchemlyb/visualisation/mbar_matrix.py @@ -0,0 +1,95 @@ +"""Functions for Plotting the overlay matrix for the MBAR estimator. + +To assess the quality of the MBAR estimator, the overlap matrix between +the lambda states can be computed and the more overlap is observed between +the states, the more reliable the estimate is. One way of accessing the +quality of the overlap matrix is by plotting it. + +The code for producing the overlap matrix plot is taken from +: `Alchemical Analysis `_. + +""" +from __future__ import division + +import matplotlib.pyplot as plt +import numpy as np + +def plot_mbar_overlap_matrix(matrix, skip_lambda_index=[], ax=None): + '''Plot the MBAR overlap matrix. + + Parameters + ---------- + matrix : numpy.matrix + DataFrame of the overlap matrix obtained from + :attr:`~alchemlyb.estimators.MBAR.overlap_matrix` + skip_lambda_index : List + list of lambda indices to be omitted from plotting process. + Default: []. + ax : matplotlib.axes.Axes + Matplotlib axes object where the plot will be drawn on. If ax=None, + a new axes will be generated. + + Returns + ------- + matplotlib.axes.Axes + An axes with the overlap matrix drawn. + + Notes + ----- + The code is taken and modified from + : `Alchemical Analysis `_ + + ''' + # Compute the size of the figure, if ax is not given. + max_prob = matrix.max() + size = len(matrix) + if ax is None: + fig, ax = plt.subplots(figsize=(size / 2, size / 2)) + ax.set_xticks([]) + ax.set_yticks([]) + ax.axis('off') + for i in range(size): + if i != 0: + ax.axvline(x=i, ls='-', lw=0.5, color='k', alpha=0.25) + ax.axhline(y=i, ls='-', lw=0.5, color='k', alpha=0.25) + for j in range(size): + if matrix[j, i] < 0.005: + ii = '' + elif matrix[j, i] > 0.995: + ii = '1.00' + else: + ii = ("{:.2f}".format(matrix[j, i])[1:]) + alf = matrix[j, i] / max_prob + ax.fill_between([i, i + 1], [size - j, size - j], [size - (j + 1), size - (j + 1)], color='k', alpha=alf) + ax.annotate(ii, xy=(i, j), xytext=(i + 0.5, size - (j + 0.5)), size=8, textcoords='data', va='center', + ha='center', color=('k' if alf < 0.5 else 'w')) + + if skip_lambda_index: + ks = [int(l) for l in skip_lambda_index] + ks = np.delete(np.arange(size + len(ks)), ks) + else: + ks = range(size) + for i in range(size): + ax.annotate(ks[i], xy=(i + 0.5, 1), xytext=(i + 0.5, size + 0.5), size=10, textcoords=('data', 'data'), + va='center', ha='center', color='k') + ax.annotate(ks[i], xy=(-0.5, size - (size - 0.5)), xytext=(-0.5, size - (i + 0.5)), size=10, textcoords=('data', 'data'), + va='center', ha='center', color='k') + ax.annotate('$\lambda$', xy=(-0.5, size - (size - 0.5)), xytext=(-0.5, size + 0.5), size=10, textcoords=('data', 'data'), + va='center', ha='center', color='k') + ax.plot([0, size], [0, 0], 'k-', lw=4.0, solid_capstyle='butt') + ax.plot([size, size], [0, size], 'k-', lw=4.0, solid_capstyle='butt') + ax.plot([0, 0], [0, size], 'k-', lw=2.0, solid_capstyle='butt') + ax.plot([0, size], [size, size], 'k-', lw=2.0, solid_capstyle='butt') + + cx = np.repeat(range(size + 1), 2) + cy = sorted(np.repeat(range(size + 1), 2), reverse=True) + ax.plot(cx[2:-1], cy[1:-2], 'k-', lw=2.0) + ax.plot(np.array(cx[2:-3]) + 1, cy[1:-4], 'k-', lw=2.0) + ax.plot(cx[1:-2], np.array(cy[:-3]) - 1, 'k-', lw=2.0) + ax.plot(cx[1:-4], np.array(cy[:-5]) - 2, 'k-', lw=2.0) + + ax.set_xlim(-1, size) + ax.set_ylim(0, size + 1) + return ax + +