A computer-implemented method for training a machine learning algorithm for predicting three-dimensional imbibition phase saturation profile for imbibition of a porous rock medium , the method comprising: deriving (S140), using a processor , a derived one-dimensional drainage phase saturation profile from a measured three-dimensional drainage phase saturation profile of a selected core plug of the porous rock medium ; calculating (S170), using the processor and a measured oil-water contact angle in the selected core plug of the porous rock medium , a plurality of synthetic values for a relative permeability of the porous rock medium and a capillary pressure of the porous rock medium ; feeding (S190) the measured three-dimensional drainage phase saturation profile , the derived one-dimensional drainage phase saturation profile , a measured one-dimensional imbibition phase saturation profile , the measured oil-water contact angle , the calculated synthetic values for the relative permeability , the calculated synthetic values for the capillary pressure , and a rock heterogeneity state into a porous media fluid flow simulator ; simulating (S200), using the porous media fluid flow simulator , a simulated three-dimensional drainage phase saturation profile and a simulated three-dimensional imbibition phase saturation profile ; calculating (S210), using the processor , a simulated one-dimensional drainage phase saturation profile from the simulated three-dimensional drainage phase saturation profile , and a simulated one-dimensional imbibition phase saturation profile from the simulated three-dimensional imbibition phase saturation profile ; and training (S220), using the simulated one-dimensional drainage phase saturation profile , the simulated three-dimensional drainage phase saturation profile , the simulated one-dimensional imbibition phase saturation profile , and the simulated three-dimensional imbibition phase saturation profile , a machine learning algorithm . deriving (S140), using a processor , a derived one-dimensional drainage phase saturation profile from a measured three-dimensional drainage phase saturation profile of a selected core plug of the porous rock medium ; calculating (S170), using the processor and a measured oil-water contact angle in the selected core plug of the porous rock medium , a plurality of synthetic values for a relative permeability of the porous rock medium and a capillary pressure of the porous rock medium ; feeding (S190) the measured three-dimensional drainage phase saturation profile , the derived one-dimensional drainage phase saturation profile , a measured one-dimensional imbibition phase saturation profile , the measured oil-water contact angle , the calculated synthetic values for the relative permeability , the calculated synthetic values for the capillary pressure , and a rock heterogeneity state into a porous media fluid flow simulator ; simulating (S200), using the porous media fluid flow simulator , a simulated three-dimensional drainage phase saturation profile and a simulated three-dimensional imbibition phase saturation profile ; calculating (S210), using the processor , a simulated one-dimensional drainage phase saturation profile from the simulated three-dimensional drainage phase saturation profile , and a simulated one-dimensional imbibition phase saturation profile from the simulated three-dimensional imbibition phase saturation profile ; and training (S220), using the simulated one-dimensional drainage phase saturation profile , the simulated three-dimensional drainage phase saturation profile , the simulated one-dimensional imbibition phase saturation profile , and the simulated three-dimensional imbibition phase saturation profile , a machine learning algorithm . <EOS>