import%20marimo%0A%0A__generated_with%20%3D%20%220.8.15%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Adef%20__()%3A%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20from%20bayesmbar%20import%20CBayesMBAR%0A%20%20%20%20return%20CBayesMBAR%2C%20np%0A%0A%0A%40app.cell%0Adef%20__(np)%3A%0A%20%20%20%20ka%2C%20kb%2C%20kc%2C%20kd%20%3D%204%2C%209%2C%2016%2C%2025%0A%20%20%20%20%23%23%20equilibrium%20positions%20of%20three%20two-dimensional%20harmonic%20oscillators%0A%20%20%20%20mu%20%3D%20%7B%0A%20%20%20%20%20%20%20%20%22a%22%3A%20np.array(%5B-1.0%2C%201.0%5D)%2C%0A%20%20%20%20%20%20%20%20%22b%22%3A%20np.array(%5B1.0%2C%201.0%5D)%2C%0A%20%20%20%20%20%20%20%20%22c%22%3A%20np.array(%5B1.0%2C%20-1.0%5D)%2C%0A%20%20%20%20%20%20%20%20%22d%22%3A%20np.array(%5B-1.0%2C%20-1.0%5D)%2C%0A%20%20%20%20%7D%0A%0A%20%20%20%20%23%23%20spring%20constants%0A%20%20%20%20k%20%3D%20%7B%0A%20%20%20%20%20%20%20%20%22a%22%3A%20np.ones(2)%20*%20ka%2C%0A%20%20%20%20%20%20%20%20%22b%22%3A%20np.ones(2)%20*%20kb%2C%0A%20%20%20%20%20%20%20%20%22c%22%3A%20np.ones(2)%20*%20kc%2C%0A%20%20%20%20%20%20%20%20%22d%22%3A%20np.ones(2)%20*%20kd%2C%0A%20%20%20%20%7D%0A%0A%20%20%20%20%23%23%20interpolate%20the%20equilibrium%20positions%20and%20spring%20constants%20between%20pairs%20of%20states%0A%20%20%20%20pairs%20%3D%20%5B(%22a%22%2C%20%22b%22)%2C%20(%22b%22%2C%20%22c%22)%2C%20(%22c%22%2C%20%22d%22)%2C%20(%22d%22%2C%20%22a%22)%2C%20(%22a%22%2C%20%22c%22)%2C%20(%22b%22%2C%20%22d%22)%5D%0A%20%20%20%20for%20s1%2C%20s2%20in%20pairs%3A%0A%20%20%20%20%20%20%20%20for%20idx%20in%20range(1%2C%204)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20mu%5B(s1%2C%20s2%2C%20idx)%5D%20%3D%20mu%5Bs1%5D%20%2B%20(mu%5Bs2%5D%20-%20mu%5Bs1%5D)%20%2F%204%20*%20idx%0A%20%20%20%20%20%20%20%20%20%20%20%20k%5B(s1%2C%20s2%2C%20idx)%5D%20%3D%20k%5Bs1%5D%20%2B%20(k%5Bs2%5D%20-%20k%5Bs1%5D)%20%2F%204%20*%20idx%0A%0A%20%20%20%20%23%23%20compute%20the%20free%20energy%20difference%20between%20pairs%20of%20states%20analytically%0A%20%20%20%20def%20k_to_sigma(k)%3A%0A%20%20%20%20%20%20%20%20return%20np.sqrt(1.0%20%2F%20k)%0A%0A%20%20%20%20deltaF_ref%20%3D%20%7B%7D%0A%20%20%20%20for%20s1%2C%20s2%20in%20pairs%3A%0A%20%20%20%20%20%20%20%20deltaF_ref%5B(s1%2C%20s2)%5D%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20-np.log(k_to_sigma(k%5Bs2%5D)).sum()%20%2B%20np.log(k_to_sigma(k%5Bs1%5D)).sum()%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20return%20deltaF_ref%2C%20idx%2C%20k%2C%20k_to_sigma%2C%20ka%2C%20kb%2C%20kc%2C%20kd%2C%20mu%2C%20pairs%2C%20s1%2C%20s2%0A%0A%0A%40app.cell%0Adef%20__(k%2C%20k_to_sigma%2C%20mu%2C%20np%2C%20pairs)%3A%0A%20%20%20%20%23%23%20sample%20configurations%20from%20all%20states%20including%20both%20endpoints%20and%20intermediates%0A%20%20%20%20n%20%3D%202000%0A%20%20%20%20np.random.seed(0)%0A%20%20%20%20x%20%3D%20%7B%7D%0A%20%20%20%20for%20s%20in%20mu.keys()%3A%0A%20%20%20%20%20%20%20%20x%5Bs%5D%20%3D%20np.random.normal(mu%5Bs%5D%2C%20k_to_sigma(k%5Bs%5D)%2C%20(n%2C%202))%0A%0A%20%20%20%20u_list%20%3D%20%5B%5D%0A%20%20%20%20for%20_s1%2C%20_s2%20in%20pairs%3A%0A%20%20%20%20%20%20%20%20key%20%3D%20%5B_s1%5D%20%2B%20%5B(_s1%2C%20_s2%2C%20idx)%20for%20idx%20in%20range(1%2C%204)%5D%20%2B%20%5B_s2%5D%0A%20%20%20%20%20%20%20%20xs%20%3D%20np.concatenate(%5Bx%5Bs%5D%20for%20s%20in%20key%5D)%0A%20%20%20%20%20%20%20%20u%20%3D%20np.stack(%0A%20%20%20%20%20%20%20%20%20%20%20%20%5Bnp.sum(0.5%20*%20k%5Bs%5D%20*%20(xs%20-%20mu%5Bs%5D)%20**%202%2C%20axis%3D1)%20for%20s%20in%20key%5D%2C%20axis%3D0%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20u_list.append(u)%0A%0A%20%20%20%20num_conf_list%20%3D%20%5B%0A%20%20%20%20%20%20%20%20%5Bu.shape%5B1%5D%20%2F%2F%20u.shape%5B0%5D%20for%20i%20in%20range(u.shape%5B0%5D)%5D%20for%20u%20in%20u_list%0A%20%20%20%20%5D%0A%20%20%20%20identical_states%20%3D%20(%0A%20%20%20%20%20%20%20%20%5B(0%2C%200)%2C%20(3%2C%204)%2C%20(4%2C%200)%5D%2C%0A%20%20%20%20%20%20%20%20%5B(0%2C%204)%2C%20(1%2C%200)%2C%20(5%2C%200)%5D%2C%0A%20%20%20%20%20%20%20%20%5B(1%2C%204)%2C%20(2%2C%200)%2C%20(4%2C%204)%5D%2C%0A%20%20%20%20%20%20%20%20%5B(2%2C%204)%2C%20(3%2C%200)%2C%20(5%2C%204)%5D%2C%0A%20%20%20%20)%0A%20%20%20%20return%20identical_states%2C%20key%2C%20n%2C%20num_conf_list%2C%20s%2C%20u%2C%20u_list%2C%20x%2C%20xs%0A%0A%0A%40app.cell%0Adef%20__(CBayesMBAR%2C%20identical_states%2C%20num_conf_list%2C%20pairs%2C%20u_list)%3A%0A%20%20%20%20cbmbar%20%3D%20CBayesMBAR(%0A%20%20%20%20%20%20%20%20u_list%2C%20num_conf_list%2C%20identical_states%2C%20random_seed%3D0%0A%20%20%20%20)%0A%0A%20%20%20%20deltaF_cbmbar%20%3D%20%7B%7D%0A%20%20%20%20for%20i%20in%20range(len(pairs))%3A%0A%20%20%20%20%20%20%20%20deltaF_cbmbar%5Bpairs%5Bi%5D%5D%20%3D%20cbmbar.DeltaF_mode%5Bi%5D%5B0%2C%20-1%5D.item()%0A%0A%20%20%20%20return%20cbmbar%2C%20deltaF_cbmbar%2C%20i%0A%0A%0A%40app.cell%0Adef%20__()%3A%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A