In the Cartesian model (model C in Figure . 1B), the positions of the hydrogen atoms in the whole system [r H1 , r H2 ] were represented directly with their Cartesian coordinates. Meanwhile, the oxygen atoms in the system (r O ) were instead represented using whitened Cartesian coordinates, as described in section II C. Taken together, the variable x C shown in Fig. , describes all the relevant N (N -1) marginal DOFs of the ice crystal. Finally, the shift and scale layer (yellow in Fig. ) was initialized to fit the one-dimensional marginal variables of x C into the working interval of the 1D splines. This step takes advantage of the minimum and maximum values explored by each 1D marginal variable in the MD training and validation data to define the one-dimensional linear transformations to carry out the scale and shift, exactly like in Ref. . Likewise, the spline coupling layers in this work were also implemented according to methods detailed in Ref. , with each marginal variable of model C assigned to a non-periodic [-1, 1] interval. The uniform log base distribution (ln p 0 ) in model C was therefore set to: