Linear ncem fixes

ncem/api/train/__init__.py

@@ -1,9 +1,17 @@
 """Initializes a train object in api."""
 import numpy as np
 
-from ncem.estimators import (Estimator, EstimatorCVAE, EstimatorCVAEncem,
-                             EstimatorDeconvolution, EstimatorED,
-                             EstimatorEDncem, EstimatorEdNcemNeighborhood,
-                             EstimatorGraph, EstimatorInteractions,
-                             EstimatorLinear, EstimatorNoGraph)
+from ncem.estimators import (
+    Estimator,
+    EstimatorCVAE,
+    EstimatorCVAEncem,
+    EstimatorDeconvolution,
+    EstimatorED,
+    EstimatorEDncem,
+    EstimatorEdNcemNeighborhood,
+    EstimatorGraph,
+    EstimatorInteractions,
+    EstimatorLinear,
+    EstimatorNoGraph,
+)
 from ncem.models import BetaScheduler

ncem/estimators/__init__.py

@@ -1,12 +1,10 @@
 """Importing estimator classes."""
-from ncem.estimators.base_estimator import (Estimator, EstimatorGraph,
-                                            EstimatorNoGraph)
+from ncem.estimators.base_estimator import Estimator, EstimatorGraph, EstimatorNoGraph
 from ncem.estimators.base_estimator_neighbors import EstimatorNeighborhood
 from ncem.estimators.estimator_cvae import EstimatorCVAE
 from ncem.estimators.estimator_cvae_ncem import EstimatorCVAEncem
 from ncem.estimators.estimator_deconvolution import EstimatorDeconvolution
 from ncem.estimators.estimator_ed import EstimatorED
-from ncem.estimators.estimator_ed_ncem import (EstimatorEDncem,
-                                               EstimatorEdNcemNeighborhood)
+from ncem.estimators.estimator_ed_ncem import EstimatorEDncem, EstimatorEdNcemNeighborhood
 from ncem.estimators.estimator_interactions import EstimatorInteractions
 from ncem.estimators.estimator_linear import EstimatorLinear

ncem/estimators/base_estimator.py

@@ -7,12 +7,20 @@
 import tensorflow as tf
 
 from ncem.utils.losses import GaussianLoss, KLLoss, NegBinLoss
-from ncem.utils.metrics import (custom_kl, custom_mae, custom_mean_sd,
-                                custom_mse, custom_mse_scaled,
-                                gaussian_reconstruction_loss, logp1_custom_mse,
-                                logp1_r_squared, logp1_r_squared_linreg,
-                                nb_reconstruction_loss, r_squared,
-                                r_squared_linreg)
+from ncem.utils.metrics import (
+    custom_kl,
+    custom_mae,
+    custom_mean_sd,
+    custom_mse,
+    custom_mse_scaled,
+    gaussian_reconstruction_loss,
+    logp1_custom_mse,
+    logp1_r_squared,
+    logp1_r_squared_linreg,
+    nb_reconstruction_loss,
+    r_squared,
+    r_squared_linreg,
+)
 
 
 def transfer_layers(model1, model2):
@@ -212,8 +220,7 @@ def _load_data(
 
         elif data_origin.startswith("cell2location_lymphnode"):
             self.targeted_assay = False
-            from ncem.data import \
-                DataLoaderCell2locationLymphnode as DataLoader
+            from ncem.data import DataLoaderCell2locationLymphnode as DataLoader
 
             self.undefined_node_types = None
 

ncem/interpretation/__init__.py

@@ -1,7 +1,9 @@
 """Importing interpretation for different model classes."""
-from ncem.interpretation.interpreter import (InterpreterCVAEncem,
-                                             InterpreterDeconvolution,
-                                             InterpreterEDncem,
-                                             InterpreterGraph,
-                                             InterpreterInteraction,
-                                             InterpreterLinear)
+from ncem.interpretation.interpreter import (
+    InterpreterCVAEncem,
+    InterpreterDeconvolution,
+    InterpreterEDncem,
+    InterpreterGraph,
+    InterpreterInteraction,
+    InterpreterLinear,
+)

ncem/models/layers/__init__.py

@@ -1,15 +1,15 @@
 """Importing custom layers for different model classes."""
 from ncem.models.layers.gnn_layers import GCNLayer, MaxLayer
-from ncem.models.layers.layer_stacks_lvm import (CondDecoder, CondEncoder,
-                                                 Decoder, Encoder,
-                                                 SamplingPrior)
-from ncem.models.layers.output_layers import (GaussianConstDispOutput,
-                                              GaussianOutput,
-                                              LinearConstDispOutput,
-                                              LinearOutput,
-                                              NegBinConstDispOutput,
-                                              NegBinOutput,
-                                              NegBinSharedDispOutput, get_out)
+from ncem.models.layers.layer_stacks_lvm import CondDecoder, CondEncoder, Decoder, Encoder, SamplingPrior
+from ncem.models.layers.output_layers import (
+    GaussianConstDispOutput,
+    GaussianOutput,
+    LinearConstDispOutput,
+    LinearOutput,
+    NegBinConstDispOutput,
+    NegBinOutput,
+    NegBinSharedDispOutput,
+    get_out,
+)
 from ncem.models.layers.preproc_input import DenseInteractions, PreprocInput
-from ncem.models.layers.single_gnn_layers import (SingleGatLayer,
-                                                  SingleLrGatLayer)
+from ncem.models.layers.single_gnn_layers import SingleGatLayer, SingleLrGatLayer

ncem/models/model_cvae.py

@@ -1,8 +1,7 @@
 import numpy as np
 import tensorflow as tf
 
-from ncem.models.layers import (Decoder, Encoder, PreprocInput, SamplingPrior,
-                                get_out)
+from ncem.models.layers import Decoder, Encoder, PreprocInput, SamplingPrior, get_out
 
 
 class ModelCVAE:

ncem/models/model_cvae_ncem.py

@@ -3,8 +3,7 @@
 import numpy as np
 import tensorflow as tf
 
-from ncem.models.layers import (CondDecoder, CondEncoder, GCNLayer, MaxLayer,
-                                PreprocInput, SamplingPrior, get_out)
+from ncem.models.layers import CondDecoder, CondEncoder, GCNLayer, MaxLayer, PreprocInput, SamplingPrior, get_out
 
 
 class ModelCVAEncem:

ncem/models/model_ed_single_ncem.py

@@ -1,10 +1,7 @@
 import tensorflow as tf
 
 from ncem.models.layers import Decoder, get_out
-from ncem.models.layers.single_gnn_layers import (SingleGatLayer,
-                                                  SingleGcnLayer,
-                                                  SingleLrGatLayer,
-                                                  SingleMaxLayer)
+from ncem.models.layers.single_gnn_layers import SingleGatLayer, SingleGcnLayer, SingleLrGatLayer, SingleMaxLayer
 
 
 class ModelEd2Ncem:

ncem/models/model_interactions.py

@@ -2,8 +2,7 @@
 
 import tensorflow as tf
 
-from ncem.models.layers import (DenseInteractions, LinearConstDispOutput,
-                                LinearOutput)
+from ncem.models.layers import DenseInteractions, LinearConstDispOutput, LinearOutput
 
 
 class ModelInteractions:

ncem/tl/__init__.py

@@ -1,10 +1,15 @@
-from ncem.tl.fit.backend.linear_model import (differential_ncem,
-                                              differential_ncem_deconvoluted,
-                                              linear_ncem,
-                                              linear_ncem_deconvoluted)
+from ncem.tl.fit.backend.linear_model import (
+    differential_ncem,
+    differential_ncem_deconvoluted,
+    linear_ncem,
+    linear_ncem_deconvoluted,
+)
 from ncem.tl.fit.backend.spline_model import (
-    get_spline_interpolation, spline_differential_ncem,
-    spline_differential_ncem_deconvoluted, spline_linear_ncem,
-    spline_linear_ncem_deconvoluted)
+    get_spline_interpolation,
+    spline_differential_ncem,
+    spline_differential_ncem_deconvoluted,
+    spline_linear_ncem,
+    spline_linear_ncem_deconvoluted,
+)
 
 from . import fit

ncem/tl/fit/backend/ablation.py

@@ -7,8 +7,7 @@
 import squidpy as sq
 from scipy.stats import linregress
 
-from ncem.tl.fit.constants import (OBS_KEY_SPLIT, OBSM_KEY_DMAT,
-                                   UNS_KEY_ABLATION, VARM_KEY_PARAMS)
+from ncem.tl.fit.constants import OBS_KEY_SPLIT, OBSM_KEY_DMAT, UNS_KEY_ABLATION, VARM_KEY_PARAMS
 from ncem.tl.fit.glm import linear_ncem
 
 

ncem/tl/fit/backend/design_matrix.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pandas as pd
 import patsy
+from formulaic import Formula
 
 from ncem.tl.fit.constants import PREFIX_INDEX, PREFIX_NEIGHBOR
 
@@ -247,8 +248,10 @@ def get_dmat_from_obs(obs: pd.DataFrame, obs_niche: pd.DataFrame, formula: str,
     obs_niche.columns = [PREFIX_NEIGHBOR + x for x in obs_niche.columns]
     # Merge sample annotation:
     obs_full = pd.concat([obs, obs_index_type, obs_niche], axis=1)
-    dmat = patsy.dmatrix(formula, obs_full)
-    dmat = pd.DataFrame(np.asarray(dmat), index=obs.index, columns=dmat.design_info.column_names)
+    columns_names = formula.replace('~0+','').split('+')
+
+    dmat = Formula(formula).get_model_matrix(obs_full)
+    dmat = pd.DataFrame(np.asarray(dmat), index=obs.index, columns=dmat.columns)
     return dmat
 
 
@@ -294,10 +297,8 @@ def get_dmats_from_deconvoluted(
         obs_index_type_x.index = obs.index
         # Merge sample annotation:
         obs_full = pd.concat([obs, obs_index_type_x, obs_niche], axis=1)
-        dmats[x] = patsy.dmatrix(formulas[x], obs_full)
-        # ensure that column names start with index type name
-        dmat_columns = [col if col.startswith(PREFIX_INDEX) else PREFIX_INDEX+x+col for col in dmats[x].design_info.column_names]
-        dmats[x] = pd.DataFrame(np.asarray(dmats[x]), index=obs.index, columns=dmat_columns)
+        dmats[x] = Formula(formulas[x]).get_model_matrix(obs_full)
+        dmats[x] = pd.DataFrame(np.asarray(dmats[x]), index=obs.index, columns=dmats[x].columns)
     return dmats
 
 
@@ -346,7 +347,8 @@ def get_dmat_global_from_deconvoluted(obs: pd.DataFrame, deconv: pd.DataFrame, f
         obs_index_type[x].index = obs.index
         # Merge sample annotation:
         obs_full = pd.concat([obs, obs_index_type_x, obs_niche], axis=1)
-        dmat_x = patsy.dmatrix(formula, obs_full)
-        dmats.append(pd.DataFrame(np.asarray(dmat_x), index=obs.index, columns=dmat_x.design_info.column_names))
+
+        dmat_x = Formula(formula).get_model_matrix(obs_full)
+        dmats.append(pd.DataFrame(np.asarray(dmat_x), index=obs.index, columns=dmat_x.columns))
     dmat = pd.concat(dmats, axis=0)
     return dmat

ncem/tl/fit/backend/linear_model.py

@@ -5,20 +5,30 @@
 import pandas as pd
 
 from ncem.tl.fit.backend.design_matrix import (
-    extend_formula_differential_ncem, extend_formula_ncem,
-    get_binary_sample_annotation_conditions, get_dmat_from_obs,
-    get_dmats_from_deconvoluted, get_obs_niche_from_graph)
+    extend_formula_differential_ncem,
+    extend_formula_ncem,
+    get_binary_sample_annotation_conditions,
+    get_dmat_from_obs,
+    get_dmats_from_deconvoluted,
+    get_obs_niche_from_graph,
+)
 from ncem.tl.fit.backend.ols_fit import ols_fit
 from ncem.tl.fit.backend.testing import test_deconvoluted, test_standard
 from ncem.tl.fit.backend.utils import write_uns
-from ncem.tl.fit.constants import (OBSM_KEY_DMAT, OBSM_KEY_DMAT_NICHE,
-                                   UNS_KEY_CELL_TYPES, UNS_KEY_CONDITIONS,
-                                   UNS_KEY_PER_INDEX, VARM_KEY_FDR_PVALS,
-                                   VARM_KEY_FDR_PVALS_DIFFERENTIAL,
-                                   VARM_KEY_PARAMS, VARM_KEY_PVALS,
-                                   VARM_KEY_PVALS_DIFFERENTIAL,
-                                   VARM_KEY_TESTED_PARAMS,
-                                   VARM_KEY_TESTED_PARAMS_DIFFERENTIAL)
+from ncem.tl.fit.constants import (
+    OBSM_KEY_DMAT,
+    OBSM_KEY_DMAT_NICHE,
+    UNS_KEY_CELL_TYPES,
+    UNS_KEY_CONDITIONS,
+    UNS_KEY_PER_INDEX,
+    VARM_KEY_FDR_PVALS,
+    VARM_KEY_FDR_PVALS_DIFFERENTIAL,
+    VARM_KEY_PARAMS,
+    VARM_KEY_PVALS,
+    VARM_KEY_PVALS_DIFFERENTIAL,
+    VARM_KEY_TESTED_PARAMS,
+    VARM_KEY_TESTED_PARAMS_DIFFERENTIAL,
+)
 
 
 def _validate_formula(formula: str, auto_keys: List[str] = []):
@@ -264,7 +274,6 @@ def linear_ncem_deconvoluted(
     )
     dmats = get_dmats_from_deconvoluted(deconv=adata.obsm[key_deconvolution], formulas=formulas, obs=adata.obs)
     for k, v in dmats.items():
-        print(k)
         dmat_key = f"{OBSM_KEY_DMAT}_{k}"
         adata.obsm[dmat_key] = v
         params = ols_fit(x_=adata.obsm[dmat_key].values, y_=adata.layers[k])

ncem/tl/fit/backend/spline_model.py

@@ -5,18 +5,26 @@
 import pandas as pd
 import patsy
 
-from ncem.tl.fit.backend.linear_model import (differential_ncem,
-                                              differential_ncem_deconvoluted,
-                                              linear_ncem,
-                                              linear_ncem_deconvoluted)
+from ncem.tl.fit.backend.linear_model import (
+    differential_ncem,
+    differential_ncem_deconvoluted,
+    linear_ncem,
+    linear_ncem_deconvoluted,
+)
 from ncem.tl.fit.backend.testing import test_deconvoluted, test_standard
 from ncem.tl.fit.backend.utils import read_uns, write_uns
-from ncem.tl.fit.constants import (PREFIX_INDEX, UNS_KEY_CELL_TYPES,
-                                   UNS_KEY_PER_INDEX, UNS_KEY_SPLINE_COEFS,
-                                   UNS_KEY_SPLINE_DF, UNS_KEY_SPLINE_FAMILY,
-                                   UNS_KEY_SPLINE_KEY_1D,
-                                   VARM_KEY_FDR_PVALS_SPLINE, VARM_KEY_PARAMS,
-                                   VARM_KEY_PVALS_SPLINE)
+from ncem.tl.fit.constants import (
+    PREFIX_INDEX,
+    UNS_KEY_CELL_TYPES,
+    UNS_KEY_PER_INDEX,
+    UNS_KEY_SPLINE_COEFS,
+    UNS_KEY_SPLINE_DF,
+    UNS_KEY_SPLINE_FAMILY,
+    UNS_KEY_SPLINE_KEY_1D,
+    VARM_KEY_FDR_PVALS_SPLINE,
+    VARM_KEY_PARAMS,
+    VARM_KEY_PVALS_SPLINE,
+)
 
 
 def get_spline_basis(df: int, key_1d_coord: str, obs: pd.DataFrame, spline_family: str):

ncem/tl/fit/backend/testing.py

@@ -127,29 +127,30 @@ def test_deconvoluted(
             for k, y in coef_to_test.items():
                 if np.all([z in parameter_names for z in y]):
                     params_idx = np.sort([parameter_names.index(z) for z in y])
-                    #theta_mle = params.values[:, params_idx]
+                    # theta_mle = params.values[:, params_idx]
                     for idx, coef in enumerate(y):
                         theta_mle = params.values[:, idx]
                         theta_sd = fisher_inv[:, idx, idx]
                         theta_sd = np.nextafter(0, np.inf, out=theta_sd, where=theta_sd < np.nextafter(0, np.inf))
                         theta_sd = np.sqrt(theta_sd)
                         theta_sd = np.expand_dims(theta_sd, axis=0)
                         theta_mle = np.expand_dims(theta_mle, axis=0)
-                        #pvals[x] = wald_test(theta_mle=theta_mle, theta_sd=theta_sd, theta0=0)
-                        #fisher_inv_subset = fisher_inv[:, idx, :][:, :, idx]
+                        # pvals[x] = wald_test(theta_mle=theta_mle, theta_sd=theta_sd, theta0=0)
+                        # fisher_inv_subset = fisher_inv[:, idx, :][:, :, idx]
                         assert coef not in pvals.keys()
                         pvals[coef] = wald_test(theta_mle=theta_mle, theta_sd=theta_sd, theta0=0)
                         qvals[coef] = correct(pvals[coef].flatten()).reshape(pvals[coef].shape)
                         tested_coefficients[coef] = theta_mle
-                        #pvals[k] = wald_test_chisq(theta_mle=theta_mle.T, theta_covar=fisher_inv_subset)
-                        #if len(idx) == 1:
+                        # pvals[k] = wald_test_chisq(theta_mle=theta_mle.T, theta_covar=fisher_inv_subset)
+                        # if len(idx) == 1:
                         #    tested_coefficients[k] = theta_mle[:, 0]
-                        #else:
+                        # else:
                         #    tested_coefficients[k] = np.zeros_like(theta_mle[:, 0]) + np.nan
         else:
             for y in coef_to_test:
                 if y in parameter_names:
                     idx = parameter_names.index(y)
+                    print(idx)
                     theta_mle = params.values[:, idx]
                     theta_sd = fisher_inv[:, idx, idx]
                     theta_sd = np.nextafter(0, np.inf, out=theta_sd, where=theta_sd < np.nextafter(0, np.inf))
@@ -163,9 +164,9 @@ def test_deconvoluted(
     qvals_arr = np.concatenate(list(qvals.values()), axis=0).T
     adata.varm[key_pval] = pd.DataFrame(pvals_arr, index=adata.var_names, columns=list(pvals.keys()))
     adata.varm[key_fdr_pval] = pd.DataFrame(qvals_arr, index=adata.var_names, columns=list(pvals.keys()))
-    #pvals_flat = np.hstack(list(pvals.values()))
-    #qvals_flat = np.hstack(list(qvals.values()))
-    #qvals_flat = correct(pvals_flat)
+    # pvals_flat = np.hstack(list(pvals.values()))
+    # qvals_flat = np.hstack(list(qvals.values()))
+    # qvals_flat = correct(pvals_flat)
     # qvals = qvals_flat.reshape((-1, len(test_keys)))
     # Write results to object:
     if key_coef is not None:

ncem/tl/fit/glm/__init__.py

@@ -1,8 +1,13 @@
-from ncem.tl.fit.backend.linear_model import (differential_ncem,
-                                              differential_ncem_deconvoluted,
-                                              linear_ncem,
-                                              linear_ncem_deconvoluted)
+from ncem.tl.fit.backend.linear_model import (
+    differential_ncem,
+    differential_ncem_deconvoluted,
+    linear_ncem,
+    linear_ncem_deconvoluted,
+)
 from ncem.tl.fit.backend.spline_model import (
-    get_spline_interpolation, spline_differential_ncem,
-    spline_differential_ncem_deconvoluted, spline_linear_ncem,
-    spline_linear_ncem_deconvoluted)
+    get_spline_interpolation,
+    spline_differential_ncem,
+    spline_differential_ncem_deconvoluted,
+    spline_linear_ncem,
+    spline_linear_ncem_deconvoluted,
+)