Mpi issues

.gitattributes

@@ -2,3 +2,4 @@
 examples/**/*.npz filter=lfs diff=lfs merge=lfs -text
 examples/**/*.npy filter=lfs diff=lfs merge=lfs -text
 examples/**/*.dat filter=lfs diff=lfs merge=lfs -text
+

.gitignore

@@ -171,6 +171,16 @@ runs_sc25/*
 settings.json
 *.out
 
+# data files that can easily be generated
+examples/brainiac/**/*.npy
+examples/brainiac/**/*.npz
+
+examples/g_ar1/**/*.npy
+examples/g_ar1/**/*.npz
+
+examples/p_ar1/**/*.npy
+examples/p_ar1/**/*.npz
+
 # Profiler outputs
 *.nsys-rep
 *.qdstrm

examples/g_ar1/generate_data.py

@@ -0,0 +1,101 @@
+import os
+import sys
+
+import numpy as np
+import scipy.sparse as sp
+from scipy.sparse.linalg import spsolve, spsolve_triangular
+from scipy.sparse import csc_matrix
+from scipy.linalg import cholesky
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+if __name__ == "__main__":
+
+    np.random.seed(5)
+    n = 1000
+
+    ## define priors
+    s2 = 5  
+    tau = 1 / s2
+    phi = 0.9  
+    # noise obs
+    obs_noise_prec = 100
+    theta_original = [
+        phi,
+        tau,
+        obs_noise_prec,
+    ]
+
+    denom = s2 * (1 - phi**2)
+
+    diag = [(1 + phi**2) / denom] * n
+    diag[0] = diag[-1] = 1 / denom
+    off_diag = [-phi / denom] * (n - 1)
+
+    Q = sp.diags([diag, off_diag, off_diag], [0, -1, 1])
+
+    # Compute sparse Cholesky factorization: Q = L @ L.T
+    # For tridiagonal matrix, we can use dense Cholesky on small blocks or scipy
+    Q_csc = Q.tocsc()
+
+    print("Q shape:", Q.shape, "Q nnz:", Q.nnz)
+    print("Q sparsity:", 100 * Q.nnz / (Q.shape[0] * Q.shape[1]), "%")
+    print(Q.toarray()[:6, :6])
+
+    # Method 1: Use dense Cholesky (for moderate sizes this is still efficient)
+    Q_dense = Q.toarray()
+    L_dense = cholesky(Q_dense, lower=True)
+    L = csc_matrix(L_dense)
+
+    print("L nnz:", L.nnz, "L sparsity:", 100 * L.nnz / (L.shape[0] * L.shape[1]), "%")
+
+    # Efficient sampling: generate z ~ N(0,I), then solve L @ u = z
+    z = np.random.normal(0, 1, size=n)
+
+    # Solve L @ u = z using sparse triangular solver
+    u = spsolve_triangular(L, z, lower=True)
+
+    # Verify the sampling worked correctly
+    print("Sample u statistics - mean:", np.mean(u), "std:", np.std(u), ". Should be around sqrt(s2) =", np.sqrt(s2))
+
+    intercept = 2
+
+    x = np.concatenate((u, [intercept]))
+    print("x: ", x[:10])
+
+    np.save("reference_outputs/x_original.npy", x)
+    np.save("inputs_ar1/x.npy", u)
+    np.save("reference_outputs/theta_original.npy", theta_original)
+
+    a_ar1 = sp.eye(n)
+    sp.save_npz("inputs_ar1/a.npz", a_ar1)
+
+    a_regression = sp.csr_matrix(np.ones((n, 1)))
+    sp.save_npz("inputs_regression/a.npz", a_regression)
+
+    eta = a_ar1 @ u + intercept
+
+    print("eta: ", eta[:6])
+    np.save("inputs_ar1/x_original.npy", eta)
+
+    noise = np.random.normal(0, np.sqrt(1 / obs_noise_prec), size=eta.shape)
+    print("noise: ", noise[:10])
+    y = eta + noise
+    np.save("y.npy", y)
+
+    print("y: ", y[:10])
+
+    Qprior = sp.block_diag([Q, sp.csr_matrix([[0.001]])])
+
+    a = sp.hstack([a_ar1, a_regression]) # a_ar1 #
+    Qcond = Qprior + obs_noise_prec * a.T @ a
+    print("Qcond: \n", Qcond.toarray()[:6,:6])
+
+    b = obs_noise_prec * a.T @ y
+    print("b: ", b[:10])
+    #x_est = np.linalg.solve(Qcond.toarray(), b)
+    x_est = spsolve(csc_matrix(Qcond), b)
+    print("norm(x - x_est): ", np.linalg.norm(x - x_est))
+
+    print("norm(eta - eta_est): ", np.linalg.norm(a @ x - a @ x_est))
+    print("normalized norm(eta - eta_est): ", np.linalg.norm(a @ x - a @ x_est) / np.linalg.norm(a @ x))

examples/g_ar1/run.py

@@ -0,0 +1,153 @@
+import sys
+import os
+
+parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.append(parent_dir)
+
+import numpy as np
+
+from dalia import xp, sp, backend_flags
+from dalia.configs import likelihood_config, dalia_config, submodels_config
+from dalia.core.model import Model
+from dalia.core.dalia import DALIA
+from dalia.submodels import AR1SubModel, RegressionSubModel
+from dalia.utils import get_host, print_msg, scaled_logit  # , extract_diagonal
+from examples_utils.parser_utils import parse_args
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+if __name__ == "__main__":
+
+    np.random.seed(3)
+
+    # load reference output
+    theta_original = np.load("reference_outputs/theta_original.npy")
+    print(
+        "theta original: ",
+        theta_original,
+    )
+
+    theta_initial = theta_original #[0.6, 1.0, 3.0]
+    print("theta initial: ", theta_initial)
+
+    x_original = np.load("reference_outputs/x_original.npy")
+    print("x original: ", x_original[:10])
+    print("dim(x original): ", x_original.shape)
+
+    ar1_dict = {
+        "type": "ar1",
+        "input_dir": f"{BASE_DIR}/inputs_ar1",
+        "phi": theta_initial[0],  # has to be between 0 and 1
+        "tau": theta_initial[1],  # assume to already be in log-scale
+        "ph_phi": {"type": "beta", "alpha": 5.0, "beta": 1.0},
+        "ph_tau": {"type": "gaussian", "mean": 0.0, "precision": 0.5},
+    }
+    ar1 = AR1SubModel(
+        config=submodels_config.parse_config(ar1_dict),
+    )
+
+    # Configurations of the regression submodel
+    regression_dict = {
+        "type": "regression",
+        "input_dir": f"{BASE_DIR}/inputs_regression",
+        "n_fixed_effects": 1,
+        "fixed_effects_prior_precision": 0.001,
+    }
+    regression = RegressionSubModel(
+        config=submodels_config.parse_config(regression_dict),
+    )
+
+    likelihood_dict = {
+        "type": "gaussian",
+        "prec_o": xp.log(theta_initial[2]),
+        # "prior_hyperparameters": {
+        #     "type": "penalized_complexity",
+        #     "alpha": 0.01,
+        #     "u": 5,
+        # },
+        "prior_hyperparameters": {"type": "gaussian", "mean": xp.log(theta_original[2]), "precision": 0.05},
+    }
+
+    model = Model(
+        submodels=[ar1, regression], #
+        likelihood_config=likelihood_config.parse_config(likelihood_dict),
+    )
+    print_msg(model)
+
+    Qprior = model.construct_Q_prior()
+    print("Qprior: \n", Qprior.toarray()[:6, :6])
+    Qinv = np.linalg.inv(Qprior.toarray())
+    geom_mean = np.exp(np.mean(np.log(Qinv.diagonal())))
+    print("Geometric mean of Qinv diagonal: ", geom_mean)
+
+    # in gaussian case x = 0, thus eta = 0
+    x_i = np.zeros(model.n_latent_parameters)
+    eta = model.a @ x_i
+    Qcond = model.construct_Q_conditional(eta=eta)
+    print("Qcond: \n", Qcond.toarray()[:6, :6])
+
+    b = model.construct_information_vector(eta=eta, x_i=x_i)
+    print("b: ", b[:10])
+
+    x_est = np.linalg.solve(Qcond.toarray(), b)
+    #print("x est: ", x_est)
+    print("norm(x_original - x_est): ", np.linalg.norm(x_original - x_est))
+
+    # Configurations of DALIA
+    dalia_dict = {
+        "solver": {"type": "dense"},
+        "minimize": {
+            "max_iter": 100,
+            "gtol": 1e-3,
+            "disp": True,
+            "maxcor": len(model.theta_external),
+        },
+        "f_reduction_tol": 1e-3,
+        "theta_reduction_tol": 1e-4,
+        "inner_iteration_max_iter": 50,
+        "eps_inner_iteration": 1e-3,
+        "eps_gradient_f": 1e-3,
+        "simulation_dir": ".",
+        "verbosity": 0,
+    }
+
+    dalia = DALIA(
+        model=model,
+        config=dalia_config.parse_config(dalia_dict),
+    )
+
+    print("initial model theta: ", model.theta_external)
+
+    print("\nCalling DALIA.run()")
+    results = dalia.run()
+
+    print_msg("\n--- Results ---")
+
+    theta = results["theta"]
+    print("theta:          ", np.round(theta, 4))
+    print("theta original: ", theta_original)
+
+    print_msg("Covariance of theta:\n", results["cov_theta"])
+    print_msg(
+        "Mean of the fixed effects:\n",
+        results["x"][-model.submodels[-1].n_fixed_effects :],
+    )
+
+    # print("x:          ", results["x"])
+    # print("x_original: ", x_original)
+
+    # print("eta: ", model.a @ x_original)
+    # print("eta est: ", model.a @ results["x"])
+
+    print_msg("\n--- Comparisons ---")
+    print("norm(eta - eta_est): ", np.linalg.norm(model.a @ x_original - model.a @ results["x"]))
+    print("normalized norm(eta - eta_est): ", np.linalg.norm(model.a @ x_original - model.a @ results["x"]) / np.linalg.norm(model.a @ x_original))
+
+    # Compare marginal variances of latent parameters
+    var_latent_params = results["marginal_variances_latent"]
+    Qconditional = dalia.model.construct_Q_conditional(eta=model.a @ model.x)
+    Qinv_ref = xp.linalg.inv(Qconditional.toarray())
+    print_msg(
+        "Norm (marg var latent - ref):    ",
+        f"{np.linalg.norm(var_latent_params - xp.diag(Qinv_ref)):.4e}",
+    )

examples/gr/preprocessing.py → examples/gr/generate_data.py

@@ -28,12 +28,12 @@
     x = L_Sigma_prior @ z
     a = sparse.random(n_observations, n_latent_parameters, density=0.5)
 
-    theta_observations = np.log(3)
+    theta_observations = 3.0
     print(f"theta_observations: {theta_observations}")
     theta_likelihood: dict = {"theta_observations": theta_observations}
 
     # generate x from a gaussian distribution of dimensions n_latent_parameters with mean 0 and precision exp(theta_observations)
-    variance = 1 / np.exp(theta_observations)
+    variance = 1 / theta_observations
     eta = a @ x
     y = np.random.normal(eta, scale=np.sqrt(variance), size=n_observations)
     print(f"x: {x}")
@@ -53,7 +53,7 @@
     sparse.save_npz(f"{path}/inputs/a.npz", a)
 
     # save original latent parameters
-    np.save(f"{path}/inputs/x_original.npy", x)
+    np.save(f"{path}/reference_outputs/x_ref.npy", x)
 
     # save original hyperparameter theta
-    np.save(f"{path}/inputs/theta_original.npy", theta_observations)
+    np.save(f"{path}/reference_outputs/theta_ref.npy", theta_observations)

examples/gr/run.py

@@ -35,8 +35,9 @@
     # Likelihood
     likelihood_dict = {
         "type": "gaussian",
-        "prec_o": 1.5,
-        "prior_hyperparameters": {"type": "gaussian", "mean": 3.5, "precision": 0.5},
+        "prec_o": 1.0,
+        "prior_hyperparameters": {"type": "gamma", "alpha": 2.0, "beta": 2.0},
+        #"prior_hyperparameters": {"type": "gaussian", "mean": 1.0, "precision": 0.5},
     }
     # Creation of the first model by combining the Regression submodel and the likelihood
     model = Model(
@@ -69,7 +70,8 @@
     results = dalia.run()
 
     print_msg("\n--- Results ---")
-    print_msg("Theta values:\n", results["theta"])
+    print_msg("Theta values external:\n", results["theta"])
+    print_msg("Theta values internal:\n", results["theta_internal"])
     print_msg("Covariance of theta:\n", results["cov_theta"])
     print_msg(
         "Mean of the fixed effects:\n",
@@ -79,6 +81,7 @@
     print_msg("\n--- Comparisons ---")
     # Compare hyperparameters
     theta_ref = xp.load(f"{BASE_DIR}/reference_outputs/theta_ref.npy")
+    print_msg("Reference theta:", theta_ref)
     print_msg(
         "Norm (theta - theta_ref):        ",
         f"{np.linalg.norm(results['theta'] - get_host(theta_ref)):.4e}",
@@ -101,7 +104,7 @@
     )
 
     # Compare marginal variances of observations
-    var_obs = dalia.get_marginal_variances_observations(theta=theta_ref, x_star=x_ref)
+    var_obs = dalia.get_marginal_variances_observations(theta_external=theta_ref, x_star=x_ref)
     var_obs_ref = extract_diagonal(model.a @ Qinv_ref @ model.a.T)
     print_msg(
         "Norm (var_obs - var_obs_ref):    ",

examples/gs_small/run.py

@@ -58,14 +58,16 @@
         likelihood_config=likelihood_config.parse_config(likelihood_dict),
     )
 
+    print_msg(model)
+
     # Configurations of DALIA
     dalia_dict = {
         "solver": {"type": "dense"},
         "minimize": {
             "max_iter": args.max_iter,
             "gtol": 1e-3,
             "disp": True,
-            "maxcor": len(model.theta),
+            "maxcor": len(model.theta_external),
         },
         "f_reduction_tol": 1e-3,
         "theta_reduction_tol": 1e-4,

examples/gst_coreg2_small/run.py

@@ -206,7 +206,7 @@
             "max_iter": args.max_iter,
             "gtol": 1e-3,
             "disp": True,
-            "maxcor": len(coreg_model.theta),
+            "maxcor": len(coreg_model.theta_external),
         },
         "f_reduction_tol": 1e-3,
         "theta_reduction_tol": 1e-4,

examples/gst_coreg3_small/run.py

@@ -266,7 +266,7 @@
             "max_iter": args.max_iter,
             "gtol": 1e-3,
             "disp": True,
-            "maxcor": len(coreg_model.theta),
+            "maxcor": len(coreg_model.theta_external),
         },
         "f_reduction_tol": 1e-3,
         "theta_reduction_tol": 1e-4,