Exhaust memory when using Eucidean cost. about moscot HOT 3 OPEN

hi @Chengwei94 ,

can you report what are the rest of the parameters for the prepare and solve method?

from moscot.

Hi, heres the snipppet I am running. I have tried both the temporal and spatial temporal problem. The joint attr is the output of a topic model where theres 50 dimension.

from jax.config import config; config.update("jax_enable_x64", True)

stp = TemporalProblem(adata)
stp = stp.score_genes_for_marginals(
gene_set_proliferation="mouse", gene_set_apoptosis="mouse"
)
adata.obs['time'] = adata.obs.timepoint.astype('str').str[1:].astype('float32').astype('category')
stp = stp.prepare(
time_key="time",
# spatial_key="spatial",
joint_attr="X_stamp",
cost = 'euclidean',
# cost = {"x": "hellinger", "y": "hellinger", "xy":"hellinger"},
device = "cpu",
# callback="local-pca",
)
stp = stp.solve(initializer="rank2", rank = 500)

Im using float64 in jax, as I encoutered an error when im using float32

from moscot.

hi @Chengwei94 sorry for late reply, yes the Euclidean cost is expected to take more memory than the squared euclidean. All the costs except the squared euclidean are less memory efficient cause they need to be factorized differently in the low rank cased, so a higher memory requirement is expected. What is the motivation to use euclidean instead of squared euclidean? and would you be able to try with a smaller rank value (e.g. 300) and on GPU? thanks!

from moscot.