Index

_

__getitem__() (dataset.CouplingsDataset method) (dataset.LinearParametrizationDataset method) (dataset.PopulationDataset method) (dataset.PopulationEvalDataset method) __init__() (dataset.CouplingsDataset method) (dataset.LinearParametrizationDataset method) (dataset.PopulationDataset method) (dataset.PopulationEvalDataset method) (models.jkonet_star.JKOnetStar method) (models.jkonet_star.JKOnetStarLinear method) (utils.density.GaussianMixtureModel method)

__len__() (dataset.CouplingsDataset method) (dataset.LinearParametrizationDataset method) (dataset.PopulationDataset method) (dataset.PopulationEvalDataset method)

A

act_fn() (networks.energies.MLP method) (networks.icnns.ICNN method)

B

beta (dataset.PopulationEvalDataset attribute) (networks.icnns.Dense attribute)

bias_init() (networks.icnns.Dense method) bohachevsky() (in module utils.functions)

C

clip_weights_icnn() (in module networks.optim) compute_couplings() (in module utils.ot) compute_couplings_sinkhorn() (in module utils.ot) compute_relevant_couplings() (in module utils.ot) const() (in module utils.features) count_parameters() (in module networks.utils) CouplingsDataset (class in dataset) create_state() (models.base.LearningDiffusionModel method) (models.jkonet.JKOnet method) (models.jkonet_star.JKOnetStar method) (models.jkonet_star.JKOnetStarLinear method) (models.jkonet_star.JKOnetStarTimePotential method)

create_state_from_params() (models.jkonet.JKOnet method) (models.jkonet_star.JKOnetStar method) (models.jkonet_star.JKOnetStarLinear method) create_train_state() (in module networks.optim) create_train_state_from_params() (in module networks.optim)

D

data_dim (dataset.PopulationEvalDataset attribute) data_generator module dataset module Dense (class in networks.icnns) densities (dataset.CouplingsDataset attribute)

densities_grads (dataset.CouplingsDataset attribute) dim_hidden (networks.energies.MLP attribute) (networks.icnns.Dense attribute) (networks.icnns.ICNN attribute) domain_from_data() (in module utils.plotting) double_exp() (in module utils.functions) dt (dataset.PopulationEvalDataset attribute) dtype (networks.icnns.Dense attribute)

E

error_wasserstein() (dataset.PopulationEvalDataset method)

error_wasserstein_cumulative() (dataset.PopulationEvalDataset method) error_wasserstein_one_step_ahead() (dataset.PopulationEvalDataset method)

F

features() (models.jkonet_star.JKOnetStarLinear method) features_dim (models.jkonet_star.JKOnetStarLinear property) features_grad() (models.jkonet_star.JKOnetStarLinear method) filename_from_args() (in module data_generator) fit() (utils.density.GaussianMixtureModel method) fixpoint_iter_bwd() (in module networks.fixpoint_loop)

fixpoint_iter_fwd() (in module networks.fixpoint_loop) flat() (in module utils.functions) flowers() (in module utils.functions) forward_sampling() (utils.sde_simulator.SDESimulator method) (utils.sde_simulator.SDESimulator_implicit_time method) friedman() (in module utils.functions) from_file() (utils.density.GaussianMixtureModel method)

G

GaussianMixtureModel (class in utils.density) generate_data_from_trajectory() (in module data_generator) get_beta() (models.base.LearningDiffusionModel method) (models.jkonet.JKOnet method) (models.jkonet_star.JKOnetStar method) (models.jkonet_star.JKOnetStarLinear method) (models.jkonet_star.JKOnetStarPotential method) get_interaction() (models.base.LearningDiffusionModel method) (models.jkonet.JKOnet method) (models.jkonet_star.JKOnetStar method) (models.jkonet_star.JKOnetStarLinear method) (models.jkonet_star.JKOnetStarPotentialInternal method) get_optimize_psi_fn() (in module models.jkonet) get_optimizer() (in module networks.optim)

get_params() (models.jkonet_star.JKOnetStar method) (models.jkonet_star.JKOnetStarLinear method) get_potential() (models.base.LearningDiffusionModel method) (models.jkonet.JKOnet method) (models.jkonet_star.JKOnetStar method) (models.jkonet_star.JKOnetStarLinear method) get_SDE_predictions() (in module utils.sde_simulator) global_norm() (in module networks.optim) gmm_density() (utils.density.GaussianMixtureModel method) gms_covs_invs (utils.density.GaussianMixtureModel attribute) gms_den (utils.density.GaussianMixtureModel attribute) gms_means (utils.density.GaussianMixtureModel attribute) gms_weights (utils.density.GaussianMixtureModel attribute) grid_from_domain() (in module utils.plotting)

H

holder_table() (in module utils.functions)

I

ICNN (class in networks.icnns) init_fn (networks.icnns.ICNN attribute) init_std (networks.icnns.ICNN attribute)

interaction (dataset.PopulationEvalDataset attribute) internal (dataset.PopulationEvalDataset attribute) ishigami() (in module utils.functions)

J

JKOnet (class in models.jkonet) JKOnetMongeGap (class in models.jkonet) JKOnetStar (class in models.jkonet_star) JKOnetStarLinear (class in models.jkonet_star)

JKOnetStarPotential (class in models.jkonet_star) JKOnetStarPotentialInternal (class in models.jkonet_star) JKOnetStarTimePotential (class in models.jkonet_star) JKOnetVanilla (class in models.jkonet)

K

kernel_init() (networks.icnns.Dense method)

L

label_mapping (dataset.PopulationEvalDataset attribute) LearningDiffusionModel (class in models.base) LinearParametrizationDataset (class in dataset) load_dataset() (models.base.LearningDiffusionModel method) (models.jkonet.JKOnet method) (models.jkonet_star.JKOnetStar method) (models.jkonet_star.JKOnetStarLinear method)

loss() (models.jkonet_star.JKOnetStar method) (models.jkonet_star.JKOnetStarPotential method) (models.jkonet_star.JKOnetStarPotentialInternal method) (models.jkonet_star.JKOnetStarTimePotential method) loss_fn_energy() (models.jkonet.JKOnet method)

M

main() (in module data_generator) (in module preprocess_rna_seq) (in module train) MLP (class in networks.energies) models.base module models.jkonet module models.jkonet_star module module data_generator dataset models.base models.jkonet models.jkonet_star networks.energies networks.fixpoint_loop networks.icnns networks.optim networks.utils preprocess_rna_seq train utils.density utils.features utils.functions utils.ot utils.plotting utils.sde_simulator

N

name (networks.energies.MLP attribute) (networks.icnns.Dense attribute) (networks.icnns.ICNN attribute) network_grad() (in module networks.utils) network_grad_time() (in module networks.utils) networks.energies module networks.fixpoint_loop module

networks.icnns module networks.optim module networks.utils module no_ground_truth (dataset.PopulationEvalDataset attribute) numpy_collate() (in module train)

O

oakley_ohagan() (in module utils.functions)

P

parent (networks.energies.MLP attribute) (networks.icnns.Dense attribute) (networks.icnns.ICNN attribute) penalize_weights_icnn() (in module networks.optim) plot_couplings() (in module utils.plotting) plot_level_curves() (in module utils.plotting) plot_predictions() (in module utils.plotting)

PopulationDataset (class in dataset) PopulationEvalDataset (class in dataset) pos_weights (networks.icnns.ICNN attribute) potential (dataset.PopulationEvalDataset attribute) precision (networks.icnns.Dense attribute) preprocess_rna_seq module

R

rbf_cubic() (in module utils.features) rbf_inverse_multiquadric() (in module utils.features) rbf_inverse_quadratic() (in module utils.features) rbf_linear() (in module utils.features)

rbf_multiquadric() (in module utils.features) rbf_quintic() (in module utils.features) rbf_thin_plate_spline() (in module utils.features) relu() (in module utils.functions) rotational() (in module utils.functions)

S

scope (networks.energies.MLP attribute) (networks.icnns.Dense attribute) (networks.icnns.ICNN attribute) SDESimulator (class in utils.sde_simulator) SDESimulator_implicit_time (class in utils.sde_simulator)

setup() (networks.energies.MLP method) (networks.icnns.ICNN method) sinkhorn_loss() (in module utils.ot) sphere() (in module utils.functions) styblinski_tang() (in module utils.functions)

T

T (dataset.PopulationEvalDataset attribute) time (dataset.CouplingsDataset attribute) to_file() (utils.density.GaussianMixtureModel method) train module train_step() (models.base.LearningDiffusionModel method) (models.jkonet.JKOnet method) (models.jkonet_star.JKOnetStar method) (models.jkonet_star.JKOnetStarLinear method)

train_test_split() (in module data_generator) trajectory (dataset.PopulationDataset attribute) (dataset.PopulationEvalDataset attribute) trajectory_only_interaction (dataset.PopulationEvalDataset attribute) trajectory_only_potential (dataset.PopulationEvalDataset attribute)

U

use_bias (networks.icnns.Dense attribute) utils.density module utils.features module utils.functions module

utils.ot module utils.plotting module utils.sde_simulator module

W

wasserstein_couplings() (in module utils.ot) wasserstein_loss() (in module utils.ot)

watershed() (in module utils.functions) wavy_plateau() (in module utils.functions) weight (dataset.CouplingsDataset attribute)

X

x (dataset.CouplingsDataset attribute)

Y

y (dataset.CouplingsDataset attribute)

Z

zigzag_ridge() (in module utils.functions)