tensorzinb-plusplus 0.0.2__py3-none-any.whl

tensorzinb/lrtest.py

@@ -0,0 +1,176 @@
+import pandas as pd
+import numpy as np
+from patsy import dmatrices
+from scipy import stats
+from .tensorzinb import TensorZINB
+
+
+class LRTest:
+    def __init__(
+        self,
+        df_data,
+        df_feature,
+        conditions,
+        nb_features,
+        nb_features_c=None,
+        infl_features=None,
+        infl_features_c=None,
+        add_intercept=True,
+        nb_only=False,
+        same_dispersion=False,
+    ):
+        self.df_data = df_data
+        self.df_feature = df_feature
+        self.conditions = list(conditions)
+        self.nb_features = list(set(nb_features) - set(conditions))
+        self.nb_features_c = nb_features_c
+        if nb_features_c is None or len(nb_features_c) == 0:
+            self._no_nb_c = True
+        else:
+            self._no_nb_c = False
+
+        if infl_features is not None and len(infl_features) > 0:
+            self.infl_features = list(set(infl_features) - set(conditions))
+            self._no_infl = False
+        else:
+            self.infl_features = None
+            self._no_infl = True
+
+        self.infl_features_c = infl_features_c
+        if infl_features_c is None or len(infl_features_c) == 0:
+            self._no_infl_c = True
+        else:
+            self._no_infl_c = False
+
+        self.nb_only = nb_only
+        if self._no_infl_c and self._no_infl:
+            self.nb_only = True
+
+        self.add_intercept = add_intercept
+        self.same_dispersion = same_dispersion
+
+        self._gen_feature_dfs()
+
+        self.res0 = None
+        self.res1 = None
+        self.df_result = None
+
+    def _get_feature_idx_map(self, Xs):
+        df_maps = []
+        for X in Xs:
+            df_f = pd.DataFrame(X.columns, columns=["feature"])
+            df_f["idx"] = range(len(df_f))
+            df_maps.append(df_f)
+        dft = pd.merge(df_maps[1], df_maps[0], on="feature")
+        idx_map = [dft.idx_x.values, dft.idx_y.values]
+        return idx_map
+
+    def _gen_feature_df(self, df_feature, features, add_intercept=True):
+        if add_intercept:
+            formula = "{} ~ {}".format("1", " + ".join(features))
+        else:
+            formula = "{} ~ {} - 1".format("1", " + ".join(features))
+
+        _, predictors = dmatrices(formula, df_feature, return_type="dataframe")
+        return predictors
+
+    def _gen_feature_dfs(self):
+        fs = [self.nb_features, self.nb_features + self.conditions]
+        Xs = []
+        dfs = []
+        for feature in fs:
+            df = self._gen_feature_df(
+                self.df_feature, feature, add_intercept=self.add_intercept
+            )
+            dfs.append(df)
+            Xs.append(df.values)
+        self.Xs = Xs
+        self.X_idx_map = self._get_feature_idx_map(dfs)
+
+        if self._no_nb_c:
+            self.X_c = None
+        else:
+            self.X_c = self._gen_feature_df(
+                self.df_feature, self.nb_features_c, add_intercept=False
+            ).values
+
+        if not self.nb_only and not self._no_infl:
+            fs = [self.infl_features, self.infl_features + self.conditions]
+            X_infls = []
+            dfs = []
+            for feature in fs:
+                df = self._gen_feature_df(
+                    self.df_feature, feature, add_intercept=self.add_intercept
+                )
+                dfs.append(df)
+                X_infls.append(df.values)
+            self.X_infls = X_infls
+            self.X_infl_idx_map = self._get_feature_idx_map(dfs)
+        else:
+            self.X_infls = [None, None]
+            self.X_infl_idx_map = None
+
+        if self._no_infl_c:
+            self.X_infl_c = None
+        else:
+            self.X_infl_c = self._gen_feature_df(
+                self.df_feature, self.infl_features_c, add_intercept=False
+            ).values
+
+    def run(self, learning_rate=0.008, epochs=5000):
+
+        zinb0 = TensorZINB(
+            self.df_data.values,
+            self.Xs[0],
+            exog_c=self.X_c,
+            exog_infl=self.X_infls[0],
+            exog_infl_c=self.X_infl_c,
+            same_dispersion=self.same_dispersion,
+            nb_only=self.nb_only,
+        )
+
+        res0 = zinb0.fit(learning_rate=learning_rate, epochs=epochs)
+        self.res0 = res0
+
+        zinb1 = TensorZINB(
+            self.df_data.values,
+            self.Xs[1],
+            exog_c=self.X_c,
+            exog_infl=self.X_infls[1],
+            exog_infl_c=self.X_infl_c,
+            same_dispersion=self.same_dispersion,
+            nb_only=self.nb_only,
+        )
+
+        weights = res0["weights"]
+        if "x_mu" in weights:
+            x_mu = np.zeros((zinb1.k_exog, zinb1.num_out))
+            x_mu[self.X_idx_map[0], :] = weights["x_mu"][self.X_idx_map[1], :]
+            weights["x_mu"] = x_mu
+
+        if "x_pi" in weights:
+            x_pi = np.zeros((zinb1.k_exog_infl, zinb1.num_out))
+            x_pi[self.X_infl_idx_map[0], :] = weights["x_pi"][self.X_infl_idx_map[1], :]
+            weights["x_pi"] = x_pi
+
+        res1 = zinb1.fit(init_weights=weights, learning_rate=learning_rate, epochs=epochs)
+        self.res1 = res1
+
+        dfr = pd.DataFrame(self.df_data.columns, columns=["test"])
+        dfr["llf0"] = res0["llfs"]
+        dfr["aic0"] = res0["aics"]
+        dfr["df0"] = res0["df"]
+        dfr["cpu_time0"] = res0["cpu_time"] / len(dfr)
+        dfr["llf1"] = res1["llfs"]
+        dfr["aic1"] = res1["aics"]
+        dfr["df1"] = res1["df"]
+        dfr["cpu_time1"] = res1["cpu_time"] / len(dfr)
+
+        dfr["llfd"] = dfr["llf1"] - dfr["llf0"]
+        dfr["aicd"] = dfr["aic1"] - dfr["aic0"]
+        dfd = dfr["df1"] - dfr["df0"]
+        dfr["pvalue"] = 1 - stats.chi2.cdf(2 * dfr["llfd"], dfd)
+
+        self.df_result = dfr
+
+        return dfr

tensorzinb/tensorzinb.py

@@ -0,0 +1,611 @@
+import warnings
+import time
+
+import contextlib
+import os
+import numpy as np
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import tensorflow as tf
+
+tf.get_logger().setLevel("ERROR")
+tf.autograph.set_verbosity(0)
+
+import tf_keras
+from tf_keras.models import Model
+from tf_keras.layers import Lambda, Input, Dense, RepeatVector, Reshape, Add
+from tf_keras.callbacks import EarlyStopping, ReduceLROnPlateau
+from tf_keras.optimizers.legacy import RMSprop
+from tf_keras import backend as K
+from scipy.special import gammaln
+import statsmodels.api as sm
+
+# Reset Keras Session
+def reset_keras():
+    K.clear_session()
+
+
+class ZINBLogLik:
+    def __init__(
+        self, pi=None, log_theta=None, nb_only=False, scope="zinb/", zero_threshold=1e-8
+    ):
+        self.pi = pi
+        self.zero_threshold = zero_threshold
+        self.scope = scope
+        self.log_theta = log_theta
+        self.nb_only = nb_only
+        self.y = None
+        self.llf = None
+
+    def loss(self, y_true, y_pred):
+        with tf.name_scope(self.scope):
+            y = tf.cast(y_true, tf.float32)
+            # mu is already in log
+            mu = tf.cast(y_pred, tf.float32)
+            log_theta = self.log_theta
+            theta = tf.math.exp(log_theta)
+
+            t1 = tf.math.lgamma(y + theta)
+            t2 = -tf.math.lgamma(theta)
+            t3 = theta * log_theta
+            t4 = y * mu
+            ty = tf.reduce_logsumexp(tf.stack([log_theta, mu], axis=0), axis=0)
+            t5 = -(theta + y) * ty
+
+            if self.nb_only:
+                result = -(t1 + t2 + t3 + t4 + t5)
+            else:
+                log_q0 = -tf.nn.softplus(-self.pi)
+                # log_q1 = -tf.nn.softplus(self.pi) = -tf.nn.softplus(-self.pi) - self.pi
+                log_q1 = log_q0 - self.pi
+
+                nb_case = -(t1 + t2 + t3 + t4 + t5 + log_q1)
+
+                p1 = theta * (log_theta - ty) + log_q1
+                zero_case = -tf.reduce_logsumexp(tf.stack([log_q0, p1], axis=0), axis=0)
+                result = tf.where(tf.less(y, self.zero_threshold), zero_case, nb_case)
+            self.llf = tf.reduce_mean(result, axis=0)
+            self.y = y
+            result = tf.reduce_sum(self.llf)
+
+        return result
+
+
+class PredictionCallback(tf_keras.callbacks.Callback):
+    def on_train_begin(self, logs={}):
+        self.weights = []
+
+    def on_epoch_end(self, epoch, logs={}):
+        ws = (
+            np.concatenate([np.array(w.flatten()) for w in self.model.get_weights()])
+        ).flatten()
+        self.weights.append(ws)
+
+
+class ReduceLROnPlateauSkip(ReduceLROnPlateau):
+    def __init__(
+        self,
+        monitor="val_loss",
+        factor=0.1,
+        patience=10,
+        verbose=0,
+        mode="auto",
+        min_delta=1e-4,
+        cooldown=0,
+        min_lr=0,
+        num_epoch_skip=3,
+        **kwargs,
+    ):
+        self.num_epoch_skip = num_epoch_skip
+
+        super(ReduceLROnPlateauSkip, self).__init__(
+            monitor=monitor,
+            factor=factor,
+            patience=patience,
+            verbose=verbose,
+            mode=mode,
+            min_delta=min_delta,
+            cooldown=cooldown,
+            min_lr=min_lr,
+            **kwargs,
+        )
+
+    def on_epoch_end(self, epoch, logs=None):
+        if epoch < self.num_epoch_skip:
+            return
+        super(ReduceLROnPlateauSkip, self).on_epoch_end(epoch, logs)
+
+
+class TensorZINB:
+    def __init__(
+        self,
+        endog,
+        exog,
+        exog_c=None,
+        exog_infl=None,
+        exog_infl_c=None,
+        same_dispersion=False,
+        nb_only=False,
+        **kwargs,
+    ):
+        self.endog = endog
+        if len(endog.shape) == 1:
+            self.num_sample = len(endog)
+            self.num_out = 1
+            self.endog = endog.reshape((-1, 1))
+        else:
+            self.num_sample, self.num_out = np.shape(endog)
+
+        df_model = 0
+        df_model_c = 0
+
+        self.exog = exog
+        df_model = np.linalg.matrix_rank(exog)
+        if len(exog.shape) == 1:
+            self.k_exog = 1
+            self.exog = exog.reshape((-1, 1))
+        else:
+            self.k_exog = exog.shape[1]
+
+        self.exog_c = exog_c
+        if exog_c is None:
+            self.k_exog_c = 0
+            self._no_exog_c = True
+            self.exog_c = np.zeros((self.num_sample, self.k_exog_c), dtype=np.float64)
+        else:
+            self.k_exog_c = exog_c.shape[1]
+            self._no_exog_c = False
+            if self.k_exog_c > 0:
+                df_model_c = df_model_c + np.linalg.matrix_rank(exog_c)
+
+        self.nb_only = nb_only
+        if exog_infl is None and exog_infl_c is None:
+            self.nb_only = True
+
+        self.exog_infl = exog_infl
+        if exog_infl is None or self.nb_only:
+            self.k_exog_infl = 0
+            self._no_exog_infl = True
+            self.exog_infl = np.ones(
+                (self.num_sample, self.k_exog_infl), dtype=np.float64
+            )
+        else:
+            self.k_exog_infl = exog_infl.shape[1]
+            self._no_exog_infl = False
+            if self.k_exog_infl > 0:
+                df_model = df_model + np.linalg.matrix_rank(exog_infl)
+
+        self.exog_infl_c = exog_infl_c
+        if exog_infl_c is None or self.nb_only:
+            self.k_exog_infl_c = 0
+            self._no_exog_infl_c = True
+            self.exog_infl_c = np.ones(
+                (self.num_sample, self.k_exog_infl_c), dtype=np.float64
+            )
+        else:
+            self.k_exog_infl_c = exog_infl_c.shape[1]
+            self._no_exog_infl_c = False
+            if self.k_exog_infl_c > 0:
+                df_model_c = df_model_c + np.linalg.matrix_rank(exog_infl_c)
+
+        df_model_each = df_model + df_model_c
+        df_model = df_model * self.num_out + df_model_c
+
+        self.same_dispersion = same_dispersion
+        if same_dispersion:
+            self.k_disperson = 1
+            df_model = df_model + 1
+        else:
+            self.k_disperson = self.num_out
+            df_model = df_model + self.num_out
+        df_model_each = df_model_each + 1
+        self.df_model = df_model
+        self.df_model_each = df_model_each
+
+        self.loglike_method = "nb2"
+
+    def fit(
+        self,
+        init_weights={},
+        init_method="poi",
+        device_type="CPU",
+        device_name=None,
+        return_history=False,
+        epochs=5000,
+        learning_rate=0.01,
+        num_epoch_skip=3,
+        is_early_stop=True,
+        is_reduce_lr=True,
+        min_delta_early_stop=0.05,
+        patience_early_stop=50,
+        factor_reduce_lr=0.8,
+        patience_reduce_lr=10,
+        min_lr=0.001,
+        reset_keras_session=False,
+        **kwargs,
+    ):
+        if device_name is None:
+            devices = tf.config.list_logical_devices(device_type)
+            device_name = devices[0].name
+
+        num_sample = self.num_sample
+        num_out = self.num_out
+        num_feat = self.k_exog
+        num_feat_infl = self.k_exog_infl
+        num_feat_c = self.k_exog_c
+        num_feat_infl_c = self.k_exog_infl_c
+        num_dispersion = self.k_disperson
+
+        # initiate weights
+        if len(init_weights) == 0:
+            if init_method == "poi":
+                init_weights = self._poisson_init()
+            elif init_method == "nb":
+                init_weights = self._nb_init()
+
+        # use distinct names to retrieve weights from layers
+        weight_keys = ["x_mu", "z_mu", "x_pi", "z_pi", "theta"]
+
+        if reset_keras_session:
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=DeprecationWarning)
+                reset_keras()
+                K.clear_session()
+
+        with tf.device(device_name):
+            tf.compat.v1.disable_eager_execution()
+
+            inputs = Input(shape=(num_feat,))
+            inputs_infl = Input(shape=(num_feat_infl,))
+            inputs_c = Input(shape=(num_feat_c,))
+            inputs_infl_c = Input(shape=(num_feat_infl_c,))
+            inputs_theta = Input(shape=(1,))
+
+            if 'x_mu' in init_weights:
+                x = Dense(num_out, use_bias=False, name='x_mu', weights=[init_weights['x_mu']])(inputs)
+            else:
+                x = Dense(num_out, use_bias=False, name='x_mu')(inputs)
+            if num_feat_c>0:
+                if 'z_mu' in init_weights:
+                    x_c = Dense(1, use_bias=False, name='z_mu', weights=[init_weights['z_mu']])(inputs_c)
+                else:
+                    x_c = Dense(1, use_bias=False, name='z_mu')(inputs_c)
+                predictions = Add()([x,x_c])
+            else:
+                predictions = x
+
+            if self.nb_only:
+                pi = None
+            else:
+                if 'x_pi' in init_weights:
+                    x_infl = Dense(num_out, use_bias=False, name='x_pi', weights=[init_weights['x_pi']])(inputs_infl)
+                else:
+                    x_infl = Dense(num_out, use_bias=False, name='x_pi')(inputs_infl)
+
+                if num_feat_infl_c>0:
+                    if 'z_pi' in init_weights:
+                        x_infl_c = Dense(1, use_bias=False, name='z_pi', weights=[init_weights['z_pi']])(inputs_infl_c)
+                    else:
+                        x_infl_c = Dense(1, use_bias=False, name='z_pi')(inputs_infl_c)
+                    pi = Add()([x_infl, x_infl_c])
+                else:
+                    pi = x_infl
+
+            if 'theta' in init_weights:
+                theta0 = Dense(num_dispersion, use_bias=False, name='theta', weights=[init_weights['theta']])(inputs_theta)
+            else:
+                theta0 = Dense(num_dispersion, use_bias=False, name='theta')(inputs_theta)
+            if self.same_dispersion:
+                theta = Reshape((num_out,))(RepeatVector(num_out)(theta0))
+            else:
+                theta = theta0
+
+            zinb = ZINBLogLik(pi, theta, nb_only=self.nb_only)
+
+            if self.nb_only:
+                output = Lambda(lambda x: x[0])([predictions, theta])
+            else:
+                output = Lambda(lambda x: x[0])([predictions, pi, theta])
+
+            model = Model(
+                inputs=[inputs, inputs_c, inputs_infl, inputs_infl_c, inputs_theta],
+                outputs=[output],
+            )
+
+            opt = RMSprop(learning_rate=learning_rate)
+            model.compile(loss=zinb.loss, optimizer=opt)
+
+            callbacks = []
+            if is_early_stop:
+                early_stop = EarlyStopping(
+                    monitor="loss",
+                    min_delta=min_delta_early_stop / num_sample,
+                    patience=patience_early_stop,
+                    mode="min",
+                )
+                callbacks.append(early_stop)
+
+            if is_reduce_lr:
+                reduce_lr = ReduceLROnPlateauSkip(
+                    monitor="loss",
+                    factor=factor_reduce_lr,
+                    patience=patience_reduce_lr,
+                    min_lr=min_lr,
+                    num_epoch_skip=num_epoch_skip,
+                )
+                callbacks.append(reduce_lr)
+
+            if return_history:
+                # this get all weights over epoch
+                get_weights = PredictionCallback()
+                callbacks.append(get_weights)
+
+            # TODO: FIX this. code randomly crashes on apple silicon M1/M2 with 
+            # error `Incompatible shapes`. code usually runs fine after second try.
+            # similar to this issue https://developer.apple.com/forums/thread/701985
+            for i in range(10):
+                try:
+                    start_time = time.time()
+                    losses = model.fit(
+                        [
+                            self.exog,
+                            self.exog_c,
+                            self.exog_infl,
+                            self.exog_infl_c,
+                            np.ones((num_sample, 1)),
+                        ],
+                        [self.endog],
+                        callbacks=callbacks,
+                        batch_size=num_sample,
+                        epochs=epochs,
+                        verbose=0,
+                    )
+                    cpu_time = time.time() - start_time
+                    break
+                except Exception as e:
+                    print(model.summary())
+                    print(e)
+                    continue
+
+            # retrieve LL
+            get_llfs = K.function(
+                [inputs, inputs_c, inputs_infl, inputs_infl_c, inputs_theta, zinb.y],
+                [zinb.llf],
+            )
+            llft = get_llfs(
+                [
+                    self.exog,
+                    self.exog_c,
+                    self.exog_infl,
+                    self.exog_infl_c,
+                    np.ones((num_sample, 1)),
+                    self.endog,
+                ]
+            )[0]
+
+            llfs = -(llft * num_sample + np.sum(gammaln(self.endog + 1), axis=0))
+            aics = -2 * (llfs - self.df_model_each)
+
+            llf = np.sum(llfs)
+            aic = -2 * (llf - self.df_model)
+
+            # get weights
+            weights = model.get_weights()
+
+            names_t = [
+                (weight.name).split("/")[0]
+                for layer in model.layers
+                for weight in layer.weights
+            ]
+            # tf layer weight has subscript in names
+            weight_names = []
+            for name in names_t:
+                matched = name
+                for nt in weight_keys:
+                    if nt in name:
+                        matched = nt
+                        break
+                weight_names.append(matched)
+
+            weights_dict = dict(zip(weight_names, weights))
+
+            res = {
+                "llf_total": llf,
+                "llfs": llfs,
+                "aic_total": aic,
+                "aics": aics,
+                "df_model_total": self.df_model,
+                "df": self.df_model_each,
+                "weights": weights_dict,
+                "cpu_time": cpu_time,
+                "num_sample": num_sample,
+                "epochs": len(losses.history["loss"]),
+            }
+
+            if return_history:
+                res["loss_history"] = losses.history["loss"]
+                res["weights_history"] = np.array(get_weights.weights)
+
+        return res
+
+    # https://github.com/statsmodels/statsmodels/blob/main/statsmodels/discrete/discrete_model.py#L3691
+    def _estimate_dispersion(self, mu, resid, df_resid=None, loglike_method="nb2"):
+        if df_resid is None:
+            df_resid = resid.shape[0]
+        if loglike_method == "nb2":
+            a = ((resid**2 / mu - 1) / mu).sum() / df_resid
+        else:  # self.loglike_method == 'nb1':
+            a = (resid**2 / mu - 1).sum() / df_resid
+        return a
+
+    def _compute_pi_init(self, nz_prob, p_nonzero, infl_prob_max=0.99):
+        ww = 1 - min(nz_prob / p_nonzero, infl_prob_max)
+        return -np.log(1 / ww - 1)
+
+    def _poisson_init_each(
+        self,
+        Y,
+        estimate_infl=True,
+        eps=1e-10,
+        maxiter=100,
+        theta_lb=0.05,
+        intercept_var_th=1e-3,
+        infl_prob_max=0.99,
+    ):
+        find_poi_sol = True
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore")
+            try:
+                poi_mod = sm.Poisson(Y, self.exog).fit(
+                    maxiter=maxiter, disp=False, warn_convergence=False
+                )
+                if np.isnan(poi_mod.params).any():
+                    find_poi_sol = False
+                else:
+                    mu = poi_mod.predict()
+                    a = self._estimate_dispersion(
+                        mu, poi_mod.resid, df_resid=poi_mod.df_resid
+                    )
+                    theta = 1 / max(a, theta_lb)
+                    x_mu = np.reshape(poi_mod.params, (self.k_exog, 1))
+            except:
+                find_poi_sol = False
+
+        if not find_poi_sol:
+            vs = np.std(self.exog, axis=0)
+            # find intercept index
+            min_idx = np.argmin(vs)
+            if vs[min_idx] < intercept_var_th:
+                x_mu = np.zeros((self.k_exog, 1))
+                mu = np.mean(Y)
+                x_mu[min_idx] = np.log(mu) / np.mean(self.exog[:, min_idx])
+                resid = Y - mu
+                a = self._estimate_dispersion(mu, resid)
+                if np.isnan(a) or np.isinf(a):
+                    a = theta_lb
+                theta = 1 / max(a, theta_lb)
+            else:
+                x_mu = None
+                return {}
+
+        weights = {"x_mu": x_mu, "theta": np.array([np.log(theta)]).reshape((-1, 1))}
+
+        if not self._no_exog_infl and estimate_infl:
+            pred = np.maximum(mu, 10 * eps)
+            p_nonzero = 1 - np.mean(np.power(theta / (theta + pred + eps), theta))
+
+            # find intercept index
+            vs = np.std(self.exog_infl, axis=0)
+            min_idx = np.argmin(vs)
+            if vs[min_idx] < intercept_var_th:
+                nz_prob = np.mean(Y > 0)
+                x_pi = np.zeros((self.k_exog_infl, 1))
+                fv = np.mean(self.exog_infl[:, min_idx])
+                w_pi = self._compute_pi_init(
+                    nz_prob, p_nonzero, infl_prob_max=infl_prob_max
+                )
+                x_pi[min_idx] = w_pi / fv
+                weights["x_pi"] = x_pi
+
+        return weights
+
+    def _poisson_init(
+        self,
+        eps=1e-10,
+        maxiter=100,
+        theta_lb=0.05,
+        intercept_var_th=1e-3,
+        infl_prob_max=0.99,
+    ):
+        x_mu = []
+        x_pi = []
+        theta = []
+        return_x_pi = True
+        return_theta = True
+        for i in range(self.num_out):
+            w = self._poisson_init_each(
+                self.endog[:, i],
+                eps=eps,
+                maxiter=maxiter,
+                theta_lb=theta_lb,
+                intercept_var_th=intercept_var_th,
+                infl_prob_max=infl_prob_max,
+            )
+
+            if len(w) == 0:
+                return {}
+
+            if "x_mu" in w:
+                x_mu.append(w["x_mu"])
+            else:
+                return {}
+
+            if "x_pi" in w:
+                x_pi.append(w["x_pi"])
+            else:
+                return_x_pi = False
+
+            if "theta" in w:
+                theta.append(w["theta"])
+            else:
+                return_theta = False
+        weights = {"x_mu": np.concatenate(x_mu, axis=1)}
+        if return_x_pi:
+            weights["x_pi"] = np.concatenate(x_pi, axis=1)
+        if return_theta:
+            t = np.concatenate(theta, axis=1)
+            if self.same_dispersion:
+                weights["theta"] = np.array(np.mean(t)).reshape((-1, 1))
+            else:
+                weights["theta"] = t
+        return weights
+
+    def _nb_init(self, infl_prob_max=0.99, intercept_var_th=1e-3):
+        nb_mod = TensorZINB(
+            self.endog,
+            self.exog,
+            exog_c=self.exog_c,
+            same_dispersion=self.same_dispersion,
+            nb_only=True,
+        )
+        nb_res = nb_mod.fit(init_method="poi")
+        weights = nb_res["weights"]
+
+        if self._no_exog_infl:
+            return weights
+        # find intercept index
+        vs = np.std(self.exog_infl, axis=0)
+        min_idx = np.argmin(vs)
+        # do not compute logit weight if there is no intercept
+        if vs[min_idx] < intercept_var_th:
+            x_pi = np.zeros((self.k_exog_infl, self.num_out))
+            fv = np.mean(self.exog_infl[:, min_idx])
+            if self.same_dispersion:
+                theta = np.exp(
+                    np.array(list(weights["theta"].flatten()) * self.num_out)
+                )
+            else:
+                theta = np.exp(weights["theta"].flatten())
+
+            mu_c = 0
+            if self.k_exog_c > 0 and "z_mu" in weights:
+                mu_c = np.dot(self.exog_c, weights["z_mu"])
+
+            for i in range(self.num_out):
+                mu = np.dot(self.exog, weights["x_mu"][:, i]) + mu_c
+                mu = np.exp(mu)
+                p_nonzero = 1 - np.mean(np.power(theta[i] / (theta[i] + mu), theta[i]))
+                nz_prob = np.mean(self.endog[:, i] > 0)
+                w_pi = self._compute_pi_init(
+                    nz_prob, p_nonzero, infl_prob_max=infl_prob_max
+                )
+                x_pi[min_idx, i] = w_pi / fv
+
+            weights["x_pi"] = x_pi
+
+            if self.k_exog_infl_c > 0:
+                weights["z_pi"] = np.zeros((self.k_exog_infl_c, 1))
+
+        return weights

tensorzinb/utils.py

@@ -0,0 +1,62 @@
+import numpy as np
+from sklearn.preprocessing import StandardScaler
+from numpy.linalg import matrix_rank
+
+def LI_vecs(dim,M):
+    LI=[M[0]]
+    idxs=[0]
+    for i in range(dim):
+        tmp=[]
+        for r in LI:
+            tmp.append(r)
+        tmp.append(M[i])                #set tmp=LI+[M[i]]
+        if matrix_rank(tmp)>len(LI):    #test if M[i] is linearly independent from all (row) vectors in LI
+            LI.append(M[i])             #note that matrix_rank does not need to take in a square matrix
+            idxs.append(i)
+    return idxs
+
+def find_independent_columns(exog):
+    if np.linalg.matrix_rank(exog)==np.shape(exog)[1]:
+        return np.array(exog.columns)
+    idxs = LI_vecs(np.shape(exog)[1], exog.T.values)
+    return np.array(exog.columns[idxs])
+
+def normalize_features(df_feature, features_to_norm):
+    scaler = StandardScaler()
+    for f in features_to_norm:
+        df_feature[f] = scaler.fit_transform(df_feature[f].values.reshape(-1, 1)).flatten()
+    return df_feature
+
+# https://stackoverflow.com/a/21739593
+def correct_pvalues_for_multiple_testing(pvalues, correction_type = "Benjamini-Hochberg"):                
+    """                                                                                                   
+    consistent with R - print correct_pvalues_for_multiple_testing([0.0, 0.01, 0.029, 0.03, 0.031, 0.05, 0.069, 0.07, 0.071, 0.09, 0.1]) 
+    """
+    from numpy import array, empty
+    pvalues = array(pvalues) 
+    n = pvalues.shape[0]                                                           
+    new_pvalues = np.zeros(n)
+    if correction_type == "Bonferroni":                                                                   
+        new_pvalues = n * pvalues
+    elif correction_type == "Bonferroni-Holm":                                                            
+        values = [ (pvalue, i) for i, pvalue in enumerate(pvalues) ]                                      
+        values.sort()
+        for rank, vals in enumerate(values):                                                              
+            pvalue, i = vals
+            new_pvalues[i] = (n-rank) * pvalue                                                            
+    elif correction_type == "Benjamini-Hochberg":                                                         
+        values = [ (pvalue, i) for i, pvalue in enumerate(pvalues) ]                                      
+        values.sort()
+        values.reverse()                                                                                  
+        new_values = []
+        for i, vals in enumerate(values):                                                                 
+            rank = n - i
+            pvalue, index = vals                                                                          
+            new_values.append((n/rank) * pvalue)                                                          
+        for i in range(0, int(n)-1):  
+            if new_values[i] < new_values[i+1]:                                                           
+                new_values[i+1] = new_values[i]                                                           
+        for i, vals in enumerate(values):
+            pvalue, index = vals
+            new_pvalues[index] = new_values[i]                                                                                                                  
+    return new_pvalues

tensorzinb_plusplus-0.0.2.dist-info/METADATA

@@ -0,0 +1,218 @@
+Metadata-Version: 2.4
+Name: tensorzinb-plusplus
+Version: 0.0.2
+Summary: Zero Inflated Negative Binomial regression for scRNA-seq (maintained fork of tensorzinb by Tao Cui & Tingting Wang)
+Project-URL: Homepage, https://github.com/saarantras/tensorzinb-plusplus
+Project-URL: Original Repository, https://github.com/wanglab-georgetown/tensorzinb
+Author: Tingting Wang
+Author-email: Tao Cui <taocui.caltech@gmail.com>
+Maintainer-email: Mackenzie Noon <me@mackenzienoon.com>
+License: Apache-2.0
+License-File: LICENSE
+Keywords: differential expression,scRNA-seq,tensorflow,zero-inflated negative binomial
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: <3.13,>=3.9
+Requires-Dist: numpy>=1.23.5
+Requires-Dist: pandas>=1.5.2
+Requires-Dist: patsy>=0.5.3
+Requires-Dist: scikit-learn>=1.2.0
+Requires-Dist: scipy>=1.9.3
+Requires-Dist: statsmodels>=0.13.5
+Requires-Dist: tensorflow>=2.16
+Requires-Dist: tf-keras>=2.16
+Description-Content-Type: text/markdown
+
+## Zero-inflated Negative Binomial Model using TensorFlow
+
+TensorZINB is a Python module that uses TensorFlow to effectively solve negative binomial (NB) and zero-inflated negative binomial (ZINB) models. One of its key strengths is its ability to accurately calculate the NB/ZINB log likelihood. Additionally, it can be used for differentially expressed gene (DEG) analysis in the context of single-cell RNA sequencing (scRNA-seq). This package distinguishes itself by ensuring numerical stability, enabling the processing of datasets in batches, and delivering superior computing speeds compared to other existing NB/ZINB solvers. To guarantee the reliability of its analysis results, TensorZINB has undergone rigorous testing against various statistical packages. TensorZINB supports the execution of various features on both the negative binomial and zero-inflated (logit) components. Furthermore, it allows for the use of common features with the same weights across multiple subjects within a batch.
+
+
+The negative binomial distribution is
+$$NB(y;\mu,\theta)=\frac{\Gamma(y+\theta)}{\Gamma(\theta)\Gamma(y+1)}\left( \frac{\theta}{\theta+\mu}\right)^\theta\left(\frac{\mu}{\theta+\mu}\right)^y$$
+where $\mu$ is the mean and $\theta$ is the dispersion parameter. For zero-inflated models, the counts are modelled as a mixture of the Bernoulli distribution and count distribution, i.e.,
+
+$$
+        Pr(Y=0)=\pi+(1-\pi)NB(0),\\
+        Pr(Y=y)=(1-\pi)NB(y),y>0.
+$$
+
+We use the following model parameterization
+
+$$
+        \log \mu_g =X_{\mu}\beta_{g,\mu}+Z_{\mu}\alpha_{\mu},
+        logit \pi_g =X_{\pi}\beta_{g,\pi}+Z_{\pi}\alpha_{\pi}, \log \theta_g = \beta_{g,\theta},
+$$
+
+where $\mu_g$ is the mean of subject $g$, $X_{\mu}$, $Z_{\mu}$, $X_{\pi}$ and $Z_{\pi}$ are feature matrices, $\beta_{g,\mu}$ and $\beta_{g,\pi}$ are coefficients for each subject $g$, $\alpha_{\mu}$ and $\alpha_{\pi}$ are common coefficients shared across all subjects.
+
+
+## Installation
+
+After downloading this repo, `cd` to the directory of downloaded repo and run:
+
+`python setup.py install`
+
+or 
+
+`pip install .`
+
+For Apple silicon (M1, M2 and etc), it is recommended to install TensorFlow by following the command in Troubleshooting Section below.
+
+## Model Estimation
+
+`TensorZINB` solves the negative binomial (NB) and zero-inflated negative binomial (ZINB) models with given read counts. 
+
+### Model initialization
+
+``` r
+TensorZINB(
+    endog,                     # counts data: number of samples x number of subjects
+    exog,                      # observed variables for the negative binomial part
+    exog_c=None,               # common observed variables across all subjects for the nb part
+    exog_infl=None,            # observed variables for the logit part
+    exog_infl_c=None,          # common observed variables across all subjects for the logit part
+    same_dispersion=False,     # whether all subjects use the same dispersion
+    nb_only=False,             # whether negative binomial only without logit or zero-inflation part
+)        
+```
+
+### Model fit
+
+``` r
+TensorZINB.fit(
+    init_weights={},          # initial model weights. If empty, init_method is used to find init weights
+    init_method="poi",        # initialization method: `poi` for Poisson and `nb` for negative binomial
+    device_type="CPU",        # device_type: `CPU` or `GPU`
+    device_name=None,         # None or one from `tf.config.list_logical_devices()`
+    return_history=False,     # whether return loss and weights history during training
+    epochs=5000,              # maximum number of epochs to run
+    learning_rate=0.008,      # start learning rate
+    num_epoch_skip=3,         # number of epochs to skip learning rate reduction
+    is_early_stop=True,       # whether use early stop
+    min_delta_early_stop=0.05,# minimum change in loss to qualify as an improvement
+    patience_early_stop=50,   # number of epochs with no improvement after which training will be stopped
+    factor_reduce_lr=0.8,     # factor by which the learning rate will be reduced
+    patience_reduce_lr=10,    # number of epochs with no improvement after which learning rate will be reduced
+    min_lr=0.001,             # lower bound on the learning rate
+    reset_keras_session=False,# reset keras session at the beginning
+)        
+```
+
+### Model results
+
+``` r
+{
+    "llf_total":              # sum of log likelihood across all subjects
+    "llfs":                   # an array contains log likelihood for each subject
+    "aic_total":              # sum of AIC across all subjects
+    "aics":                   # an array contains AIC for each subject
+    "df_model_total":         # total degree of freedom of all subjects
+    "df_model":               # degree of freedom for each subject
+    "weights":                # model weights
+    "cpu_time":               # total computing time for all subjects  
+    "num_sample":             # number of samples
+    "epochs":                 # number of epochs run
+    "loss_history":           # loss history over epochs if return_history=True
+    "weights_history":        # weights history over epochs if return_history=True
+}     
+```
+
+## DEG Analysis
+
+`LRTest` provides utility for scRNA-seq DEG analysis. It runs the likelihood ratio test (LRT) by computing the log likelihood difference with and without conditions being added to the model.
+
+To construct a `LRTest` object, we use
+``` r
+LRTest(
+    df_data,                 # count data frame. columns: subjects (genes), rows: samples
+    df_feature,              # feature data frame. columns: features, rows: samples
+    conditions,              # list of features to test DEG, e.g., diagnosis
+    nb_features,             # list of features for the negative binomial model
+    nb_features_c=None,      # list of common features for the negative binomial model
+    infl_features=None,      # list of features for the zero inflated (logit) model
+    infl_features_c=None,    # list of common features for the zero inflated (logit) model
+    add_intercept=True,      # whether add intercept. False if df_feature already contains intercept
+    nb_only=False,           # whether only do negative binomial without zero inflation
+    same_dispersion=False,   # whether all subjects use the same dispersion
+)        
+```
+
+We then call `LRTest.run` to run the likelihood ratio test
+``` r
+LRTest.run(
+    learning_rate=0.008,     # learning rate
+    epochs=5000,             # number of epochs run
+)        
+```
+
+The `LRTest.run` returns a result dataframe `dfr` with columns:
+``` r
+[
+    "ll0":                   # log likelihood without conditions
+    "aic0":                  # AIC without conditions
+    "df0":                   # degree of freedom without conditions
+    "cpu_time0":             # computing time for each subject without conditions
+    "ll1":                   # log likelihood without conditions
+    "aic1":                  # AIC with conditions
+    "df1":                   # degree of freedom with conditions
+    "cpu_time1":             # computing time for each subject with conditions
+    "lld":                   # ll1 - ll0
+    "aicd":                  # aic1 - aic0
+    "pvalue":                # p-value: 1 - stats.chi2.cdf(2 * lld, df1 - df0)
+]      
+```
+
+
+`tensorzinb.utils` provides utility functions:
+
+- `normalize_features`: normalize scRNA-seq features by removing the mean and scaling to unit variance.
+- `correct_pvalues_for_multiple_testing`: correct pvalues for multiple testing in Python, which is the same as `p.adjust` in `R`.
+
+We can further correct pvalues for multiple testing by calling `correct_pvalues_for_multiple_testing(dfr['pvalue'])`.
+ 
+## Example
+
+An example code to show how to use `TensorZINB` and `LRTest` to perform DEG analysis can be found at [`examples/deg_example.ipynb`](examples/deg_example.ipynb). The example runs DEG analysis on a sample dataset with 17 clusters and 20 genes in each cluster. 
+
+
+## Tests
+
+In `tests/tensorzinb.ipynb`, we show several tests:
+
+- validate the Poisson weights initialization.
+- compare with `statsmodels` for negative binomial model only without zero-inflation to make sure the results match.
+- show `statsmodels` is not numerically stable for zero-inflated negative binomial. `statsmodels` can only return results when initialized with TensorZINB results. TensorZINB results match the true parameters used to generate the samples.
+
+More tests can be found in https://github.com/wanglab-georgetown/countmodels/blob/main/tests/zinb_test.ipynb
+
+
+## Troubleshooting
+
+### Run on Apple silicon
+To run tensorflow on Apple silicon (M1, M2, etc), install TensorFlow using the following:
+
+`conda install -c apple tensorflow-deps`
+
+`python -m pip install tensorflow-macos==2.9.2`
+
+`python -m pip install tensorflow-metal==0.5.1`
+
+### Feature normalization
+
+If the solver cannot return correct results, please ensure features in $X$ are normalized by using `StandardScaler()`. Please refer to the example in [`examples/deg_example.ipynb`](examples/deg_example.ipynb).
+
+
+## Reference
+Cui, T., Wang, T. [A Comprehensive Assessment of Hurdle and Zero-inflated Models for Single Cell RNA-sequencing Analysis](https://doi.org/10.1093/bib/bbad272), Briefings in Bioinformatics, July 2023. https://doi.org/10.1093/bib/bbad272
+
+## Support and Contribution
+If you encounter any bugs while using the code, please don't hesitate to create an issue on GitHub here.

tensorzinb_plusplus-0.0.2.dist-info/RECORD

@@ -0,0 +1,8 @@
+tensorzinb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tensorzinb/lrtest.py,sha256=eAwPdr5cKdJzVED5IGdHe9ZTjewhzqu_56Q6HSqGrt4,5680
+tensorzinb/tensorzinb.py,sha256=mR995dyAegF51pIcDf_WtLsxAVTUn6A6ycMTTT9q6cQ,20535
+tensorzinb/utils.py,sha256=whRWe4TpGqsQ3R6p8eHjvquEBDm0IKX00WMisduN1sg,3429
+tensorzinb_plusplus-0.0.2.dist-info/METADATA,sha256=AMuET10jEXK_oELgMvRLFRrHN0tv3joAKQN2NgArprI,10850
+tensorzinb_plusplus-0.0.2.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+tensorzinb_plusplus-0.0.2.dist-info/licenses/LICENSE,sha256=vdzedejlUZ0EzkJtQR3a_QIOkx5EHPrn04D_T5k-uDk,566
+tensorzinb_plusplus-0.0.2.dist-info/RECORD,,

tensorzinb_plusplus-0.0.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

tensorzinb_plusplus-0.0.2.dist-info/licenses/LICENSE

@@ -0,0 +1,13 @@
+Copyright 2020-2023 Tao Cui, Tingting Wang
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.