Transition from computeExpectations to compute_expectations - NxK instead of NxKxK matrices

[cgaimoli]

Dear community,
I am revisiting my old scripts (https://github.com/choderalab/pymbar-datasets/blob/master/gas-properties), so they can work with Pymbar v4.0. As I am starting to investigate the differences, I noticed that one error I am getting is the same that occurs when I try to run the harmonic-oscillators.py example with observe = "RMS displacement". The error is the following:

WARNING:pymbar.mbar:dim=3 for (state_dependent==True) matrices for observables and dim=2 for (state_dependent==False) observables are deprecated; we suggest you convert to NxK form instead of NxKxK form.
Traceback (most recent call last):
  File "/harmonic-oscillators.py", line 183, in <module>
    results = mbar.compute_expectations(A_kn)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/python3.12/site-packages/pymbar/mbar.py", line 1251, in compute_expectations
    inner_results = self.compute_expectations_inner(
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/python3.12/site-packages/pymbar/mbar.py", line 935, in compute_expectations_inner
    Log_W_nk[:, sa] = np.log(A_n[i, ri]) + Log_W_nk[:, l]
                             ~~~^^^^^^^
IndexError: index 6 is out of bounds for axis 1 with size 6

Could anyone provide some guidance on how to tackle this? Since the RMS_displacement is inherently a 3D matrix, I could not find a solution to convert it to a 2D array as required to run results = mbar.compute_expectations(A_kn). Any help would be much appreciated!

[mrshirts]

So, you observables are fundamentally a 1d matrix - even if you are looking at an observable for 3 dimensions x N molecules, that's a 1D observable of length 3N. The second dimension is if that observable differs depending on the state. The energy changes depending on the state - observables depending only on the coordinates do not depend on state.