Exhaust memory when using Eucidean cost.

[Chengwei94]

Hi there,

I notice that I encoutered the exhaust memory error when using the Euclidean cost, but not when Im using SqEuclidean cost. I am running on MOSTA and it fails on the E12.5 to E13.5, (50k - 70k celll). Is there a reason for why the Euclidean cost matrix is less efficient?

[giovp]

hi @Chengwei94 ,

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

[Chengwei94]

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

[giovp]

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!