fix!(mace): send forces in C order expected by Fortran

[FelyCZ]

hello,
I was informed by Jakub Vokrouhlicky, that during his experiments with alanine, all the simulations using MACE potential were ending with an explosion.

investigation of the cause confirmed that there is a critical bug in the way forces are transferred between python server and fortran. the bug was missed during reviews due to an unfortunate selection of test molecule and a bug in HarmonicModel.

---

the bug

the fortran expects matrix in (3, natom) column-major ordered array $\big[F_{x,0}, F_{y,0}, F_{z,0}, F_{x,1}, F_{y,1}, F_{z,1} \dots\big]$, however, calling forces.T.copy() creates an array of shape (3, natom) in row-major order (which is default for C lang - CPython) $\big[F_{x,0}, F_{x,1},\dots, F_{y,0}, F_{y,1}, \dots, F_{z,0}, F_{z,1} \dots\big]$.

ABIN/interfaces/MACE/mace_server.py

Lines 303 to 309 in 9c4c611

	# Send forces (3*natom doubles, in Hartree/Bohr)
	# Transpose back to (3, natom) to match Fortran column-major layout
	if forces.dtype != np.float64:
	forces_send = forces.T.astype(np.float64)
	else:
	forces_send = forces.T.copy()
	abin_comm.Send([forces_send, MPI.DOUBLE], dest=0, tag=MACE_TAG_DATA)

why test passed

that's a combination of two problems.

1. harmonic model has accidental transposition
   fortran sends (3, natom) in column-major order, which means (natom, 3) in python's row-major order. calling .T makes it (3, natom) instead. another transposition is done in main().
   

   ABIN/interfaces/MACE/mace_server.py

   Line 145 in 9c4c611

   coords_t = coords_bohr.T # (natom, 3)

2. water molecule was used for testing
   3 atoms & 3 coords. this means the matrix is a square and therefore all subsequent calculation run without problems. the second transposition in main() then canceled out the wrong order of send values.

the fix

1. removed accidental coords_bohr.T transposition in HarmonicModel
2. removed the forces.T transposition in mace_server.py to match to format of array expected by fortran
3. replaced water for alanine in test to prevent any potential square matrix cancellation bugs in the future

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.02%. Comparing base (9c4c611) to head (e128704).

Additional details and impacted files

@@           Coverage Diff           @@
##           master     #259   +/-   ##
=======================================
  Coverage   84.02%   84.02%           
=======================================
  Files          49       49           
  Lines        6787     6787           
  Branches      786      786           
=======================================
  Hits         5703     5703           
  Misses        878      878           
  Partials      206      206           

Flag	Coverage Δ
unittests	18.25% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[danielhollas]

Hey, thanks for the PR!

Can I just double check, was Jakub using the latest ABIN version, i.e. the current master branch? Because I believe I have already fixed this bug when I was updating and testing your original PR, since we were also seeing a blow up when simulating our system (methylhydroperoxide molecule).

[FelyCZ]

according to him, he was. and I actually rerun his experiments on my machine with the latest version as well.
I even tried to recompile abin from scratch (the master branch) and run it again, but i still get the same problem.

could you give it a try yourself?

input.in

&general
pot='_mace_'	! where do we obtain forces?
irest=0,		! Should we restart from restart.xyz? (ignoring mini.dat)

mdtype='md',	! md = classical MD
nwalk=1,		! number of random walkers 
nstep=100,		! number of steps
dt=10.,			! timestep [au]


nwrite=1,		! how often some output should be printed (estimators etc.)
nwritex=5,		! how often should we print coordinates?
nrest=10,		! how often we print restart files?
nwritev=5,		! how often we print velocities?
nwritef=5,		! how often we print forces?
/

! Thermostat options
&thermostat
temp=298.15,		! temperature (Kelvins) for initial Maxwell-Boltzmann sampling of velocities
therm='none',
/

mini.xyz

13
Alanine
N    1.112912   -1.326324   -0.010213
C    0.646282    0.026292   -0.379256
C    1.442610    1.176654    0.244491
C   -0.863016    0.147188   -0.057389
O   -1.462280   -1.024564    0.210715
O   -1.459630    1.200067   -0.077965
H    1.010237    2.139592   -0.039738
H    2.486373    1.142163   -0.087301
H    1.423939    1.112300    1.339928
H    0.706746    0.101825   -1.472948
H    1.803065   -1.683981   -0.665144
H    1.541319   -1.329334    0.914434
H   -0.742040   -1.703133    0.173186

i used MACE-OFF23_medium model but according to Kuba it happens with different models as well.
the mace_server.py is without any modification and run script has just the environment variables adjusted for my machine.

[danielhollas]

Thank you, I'll test (probably next week since I am out of office). We're using our own trained model so that might be the reason for the different behaviour that we're seeing somehow.

[danielhollas]

Sorry for the delay, I've finally tested things and I indeed see a blowup with the MACE-OFF23_medium model on master branch, and correct dynamics on this PR.

On the other hand, our custom MACE model for the methylhydroperoxide molecule is working on both the master branch and this branch, and I don't yet understand why exactly. The difference seems to be that our model is in dtype=float32 wherease the MACE-OFF23 is in dtype64. Need to think a bit more how that could possibly influence the shape / ordering of things...